"""iLocal-style context extraction for source code.

Adapts Bex 2007's iLocal algorithm from YAML/XML to source code.
Instead of hard-coding directory as the grouping key, this module
extracts (context, sequence) pairs and lets reduce find natural groupings.

Context = structural position that determines the call sequence.
For YAML: context = YAML key.
For source code: context = file path component, class name, import set, etc.
"""

import os
from collections import Counter, defaultdict


def extract_contexts_by_dir(sequences, seq_files, project_root):
    """Context = relative directory path (like current package grouping).

    Returns dict: {context_key: [original_sequence, ...]}
    """
    contexts = defaultdict(list)
    for seq, fp in zip(sequences, seq_files):
        rel = os.path.relpath(os.path.dirname(fp), project_root)
        if rel == ".":
            rel = ""
        contexts[rel].append(seq)
    return dict(contexts)


def extract_contexts_by_file(sequences, seq_files, project_root):
    """Context = individual file path.

    Returns dict: {context_key: [original_sequence, ...]}
    """
    contexts = defaultdict(list)
    for seq, fp in zip(sequences, seq_files):
        rel = os.path.relpath(fp, project_root)
        contexts[rel].append(seq)
    return dict(contexts)


def extract_contexts_by_parent_dir(sequences, seq_files, project_root):
    """Context = parent of directory (one level coarser than dir).

    Returns dict: {context_key: [original_sequence, ...]}
    """
    contexts = defaultdict(list)
    for seq, fp in zip(sequences, seq_files):
        rel_dir = os.path.relpath(os.path.dirname(fp), project_root)
        if rel_dir == ".":
            rel_dir = ""
        parts = rel_dir.replace(os.sep, "/").rstrip("/").split("/")
        parent = "/".join(parts[:-1]) if len(parts) > 1 else ""
        contexts[parent].append(seq)
    return dict(contexts)


def extract_contexts_by_file_depth(sequences, seq_files, project_root, depth=2):
    """Context = first N components of file path.

    Returns dict: {context_key: [original_sequence, ...]}
    """
    contexts = defaultdict(list)
    for seq, fp in zip(sequences, seq_files):
        rel = os.path.relpath(os.path.dirname(fp), project_root)
        if rel == ".":
            rel = ""
        parts = rel.replace(os.sep, "/").rstrip("/").split("/")
        key = "/".join(parts[:depth]) if parts and parts[0] else ""
        contexts[key].append(seq)
    return dict(contexts)


def extract_contexts_by_import_set(sequences, seq_files, project_root):
    """Context = unique import set of the file.

    Methods from files with identical imports share a context.
    """
    contexts = defaultdict(list)
    for seq, fp in zip(sequences, seq_files):
        imports = _file_import_signature(fp)
        contexts[imports].append(seq)
    return dict(contexts)


def _file_import_signature(fp):
    """Build a canonical import signature for a file."""
    try:
        with open(fp) as f:
            lines = [f.readline() for _ in range(100)]
    except OSError:
        return ""
    imports = []
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            imports.append(stripped)
    return "|".join(sorted(set(imports)))


def extract_contexts_by_symbol_overlap(sequences, seq_files, project_root, min_overlap=0.5):
    """Context = cluster of methods that share >= min_overlap of their symbols.

    Uses union-find for transitive clustering.
    """
    symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]
    n = len(symbol_seqs)
    if n == 0:
        return {}

    # Build symbol sets
    sym_sets = [set(s) for s in symbol_seqs]

    # Union-find
    parent = list(range(n))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(x, y):
        px, py = find(x), find(y)
        if px != py:
            parent[px] = py

    # Cluster by symbol overlap
    for i in range(n):
        for j in range(i + 1, n):
            if not sym_sets[i] or not sym_sets[j]:
                continue
            overlap = len(sym_sets[i] & sym_sets[j]) / max(len(sym_sets[i] | sym_sets[j]), 1)
            if overlap >= min_overlap:
                union(i, j)

    # Group by cluster — keep original tuples
    clusters = defaultdict(list)
    for i in range(n):
        root = find(i)
        clusters[root].append(sequences[i])

    return {f"cluster_{k}": v for k, v in clusters.items()}


def reduce_contexts(context_groups):
    """iLocal reduce: merge contexts with identical signatures.

    Signature = sorted set of (length, first_symbol, last_symbol) across
    all sequences in the context.

    Args:
        context_groups: dict of {context_key: [original_sequence, ...]}

    Returns:
        dict: {generalized_context: [original_sequence, ...]}
    """
    if not context_groups:
        return {}

    signature_map = {}
    for ctx, seqs in context_groups.items():
        sig_parts = []
        for s in seqs:
            symbol_seq = [text for _, text, _ in s] if s and isinstance(s[0], tuple) else s
            first = symbol_seq[0] if symbol_seq else "∅"
            last = symbol_seq[-1] if symbol_seq else "∅"
            sig_parts.append((len(symbol_seq), first, last))
        signature = tuple(sorted(set(sig_parts)))
        if signature not in signature_map:
            signature_map[signature] = []
        signature_map[signature].append(ctx)

    result = {}
    for sig, ctx_list in signature_map.items():
        merged_ctx = "|".join(sorted(ctx_list)[:3])  # cap label length
        if len(ctx_list) > 3:
            merged_ctx += f"|+{len(ctx_list)-3}"
        merged_seqs = []
        for ctx in ctx_list:
            merged_seqs.extend(context_groups[ctx])
        result[merged_ctx] = merged_seqs

    return result


# Context extraction strategies registry
CONTEXT_STRATEGIES = {
    "dir": extract_contexts_by_dir,
    "file": extract_contexts_by_file,
    "parent_dir": extract_contexts_by_parent_dir,
    "depth_2": lambda s, f, p: extract_contexts_by_file_depth(s, f, p, depth=2),
    "depth_3": lambda s, f, p: extract_contexts_by_file_depth(s, f, p, depth=3),
    "imports": extract_contexts_by_import_set,
    "symbol_overlap": extract_contexts_by_symbol_overlap,
}
