"""Orchestrator: directory scan → preprocess → frequency filter → ensemble infer.

Usage:
    python -m bex.tag_preprocessor.analyze <directory> [options]

Runs the full pipeline over a directory of source files.
"""

import argparse
import json
import os
import re
import sys
import time
from pathlib import Path
from collections import Counter
from concurrent.futures import ProcessPoolExecutor, as_completed

import pathspec

from .code import preprocess_by_method, extract_arg_info, _summarize_arg_info
from bex.ensemble import infer_ensemble
from bex.gbnf import validate_sore, grammar_structure_score

SUPPORTED_EXTENSIONS = {
    ".py", ".js", ".ts", ".kt", ".rb", ".go", ".rs", ".java", ".c", ".cpp",
}

FALLBACK_SKIP = {
    ".git", "node_modules", "build", "dist", "target", "bin", "obj",
    "__pycache__", ".gradle", ".mvn", ".idea", ".vscode", ".tox",
    ".venv", "venv", "env", ".env", "out",
}

TEST_PATH_PATTERNS = [
    "**/test/**", "**/tests/**", "**/*Test*/**",
    "**/*_test.*", "**/*_spec.*", "**/test_*.py",
    "**/*Test.kt", "**/*Test.java", "**/*Test.js",
    "**/*Spec.*", "**/*_test.go",
]


def _is_main_source(filepath):
    """Return True if file path looks like main source (not test)."""
    for pattern in TEST_PATH_PATTERNS:
        if _match_glob(filepath, pattern):
            return False
    return True


def _load_gitignore(dir_path):
    """Load .gitignore from dir_path root, return PathSpec or None."""
    path = os.path.join(dir_path, ".gitignore")
    if os.path.isfile(path):
        with open(path) as f:
            return pathspec.PathSpec.from_lines("gitwildmatch", f)
    return None


def _match_glob(filepath, pattern):
    """Match filepath against a gitignore-style glob pattern.

    Uses pathspec for proper **/ recursion support.
    """
    spec = pathspec.PathSpec.from_lines("gitwildmatch", [pattern])
    return spec.match_file(filepath)


IMPORT_PATTERNS = [
    re.compile(r"^\s*import\s+"),
    re.compile(r"^\s*from\s+"),
    re.compile(r"^\s*require_relative\s+"),
    re.compile(r"^\s*require\s+"),
    re.compile(r"^\s*#\s*include\s+"),
    re.compile(r"^\s*use\s+"),
    re.compile(r"^\s*include\s+"),
]


def _file_to_package(fp, project_root):
    """Infer the package/module namespace from a file path.

    Uses the directory of the file relative to the project root.
    Zero per-language or per-convention assumptions — pure path math.
    """
    rel = os.path.relpath(os.path.dirname(fp), project_root)
    if rel == ".":
        return ""
    return rel


def _top_packages(file_paths, project_root, top_n=3):
    """Return the most common packages among a set of files."""
    pkg_counts = Counter()
    for fp in file_paths:
        pkg = _file_to_package(fp, project_root)
        pkg_counts[pkg] += 1
    return [pkg for pkg, _ in pkg_counts.most_common(top_n)]


def _extract_imports(file_paths):
    """Extract unique import lines from source files.

    Scans top 200 lines of each file for common import patterns
    across all supported languages. Deduplicates across files.
    """
    seen = set()
    result = []
    for fp in sorted(file_paths):
        try:
            with open(fp) as f:
                for i, line in enumerate(f):
                    if i >= 200:
                        break
                    stripped = line.strip()
                    if any(p.match(stripped) for p in IMPORT_PATTERNS):
                        if stripped not in seen:
                            seen.add(stripped)
                            result.append(stripped)
        except OSError:
            continue
    return result


def _build_arg_patterns(file_paths):
    """Extract merged argument patterns across a set of files."""
    merged = {}
    for fp in file_paths:
        try:
            with open(fp) as f:
                code = f.read()
        except OSError:
            continue
        info = extract_arg_info(fp, code)
        for call_name, observations in info.items():
            merged.setdefault(call_name, []).extend(observations)
    return _summarize_arg_info(merged)


def scan_directory(dir_path, gitignore_spec=None):
    """Walk dir_path, return dict mapping extension → [file paths].

    Uses .gitignore patterns from the target directory to skip
    ignored files/dirs. Falls back to FALLBACK_SKIP for common
    build/dependency dirs when no .gitignore exists.
    """
    if gitignore_spec is None:
        gitignore_spec = _load_gitignore(dir_path)

    result = {}
    for root, dirs, files in os.walk(dir_path):
        rel_root = os.path.relpath(root, dir_path)
        if rel_root == ".":
            rel_root = ""

        pruned = []
        for d in dirs:
            rel_dir = os.path.join(rel_root, d) if rel_root else d
            if gitignore_spec and gitignore_spec.match_file(rel_dir):
                continue
            if gitignore_spec is None and d in FALLBACK_SKIP:
                continue
            pruned.append(d)
        dirs[:] = pruned

        for f in files:
            rel_file = os.path.join(rel_root, f) if rel_root else f
            if gitignore_spec and gitignore_spec.match_file(rel_file):
                continue
            ext = os.path.splitext(f)[1].lower()
            if ext in SUPPORTED_EXTENSIONS:
                result.setdefault(ext, []).append(os.path.join(root, f))
    return result


DEFAULT_COVERAGE = 0.05

VERBOSE = False
_vstart = 0.0


def _vprint(*args, **kwargs):
    if VERBOSE:
        elapsed = time.time() - _vstart
        print(f"[{elapsed:6.1f}s]", *args, file=sys.stderr, **kwargs)


def frequency_filter(sequences, min_coverage=0.2):
    """Remove symbols appearing in fewer than min_coverage fraction of files.

    Args:
        sequences: list of lists of (capture_name, text, line_number) tuples.
        min_coverage: minimum fraction of files a symbol must appear in.

    Returns:
        Filtered sequences with rare symbols removed.
    """
    if not sequences:
        return sequences

    n_files = len(sequences)
    threshold = max(1, int(n_files * min_coverage))

    symbol_file_count = Counter()
    for seq in sequences:
        seen = set()
        for _, text, _ in seq:
            if text not in seen:
                symbol_file_count[text] += 1
                seen.add(text)

    keep = {text for text, count in symbol_file_count.items()
            if count >= threshold}

    filtered = []
    for seq in sequences:
        new_seq = [(cap, text, line) for cap, text, line in seq if text in keep]
        filtered.append(new_seq)

    return filtered


def _preprocess_file(fp):
    """Preprocess one file. Module-level for ProcessPoolExecutor."""
    with open(fp) as f:
        code = f.read()
    sequences = []
    for method_seq in preprocess_by_method(fp, code):
        if method_seq:
            sequences.append(method_seq)
    return (fp, sequences)


def _preprocess_files(file_paths):
    """Preprocess multiple files in parallel."""
    sequences = []
    seq_files = []
    n_workers = os.cpu_count()
    with ProcessPoolExecutor(max_workers=n_workers) as ex:
        futures = {ex.submit(_preprocess_file, fp): fp for fp in file_paths}
        for f in as_completed(futures):
            fp, method_seqs = f.result()
            for seq in method_seqs:
                sequences.append(seq)
                seq_files.append(fp)
    return sequences, seq_files


def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False, method='langsize', min_methods=3):
    """Run full pipeline: preprocess → frequency filter → ensemble infer.

    Returns:
        list of (label, ensemble_result_dict, method_count, meta) tuples.
        meta = {"files": set(paths), "imports": [lines], "arg_patterns": {...}, "packages": [...]}.
    """
    sequences, seq_files = _preprocess_files(file_paths)
    if not sequences:
        return []

    sequences = frequency_filter(sequences, min_coverage=min_coverage)

    cluster_fps = set(seq_files)
    imports = _extract_imports(cluster_fps)
    arg_patterns = _build_arg_patterns(cluster_fps)
    packages = _top_packages(cluster_fps, project_root)

    symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]

    # Diversity threshold
    n_methods = len(symbol_seqs)
    meta = {"files": cluster_fps, "imports": imports, "arg_patterns": arg_patterns, "packages": packages}
    if n_methods < min_methods:
        meta["skip_reason"] = "too_few_methods"
        return [("(all methods)", None, len(sequences), meta)]
    unique_seqs = len(set(tuple(s) for s in symbol_seqs))
    unique_ratio = unique_seqs / n_methods
    if unique_ratio > 0.9:
        meta["skip_reason"] = "too_diverse"
        return [("(all methods)", None, len(sequences), meta)]

    result = infer_ensemble(symbol_seqs, kmax=kmax, N=N, prefer=prefer, min_coverage=min_coverage, include_kore=include_kore, method=method)

    return [("(all methods)", result, len(sequences), meta)]


def _split_by_first_symbol(symbol_seqs, min_subgroup=3):
    """Split sequences by first symbol to separate mixed patterns.

    Returns:
        dict mapping first_symbol → list of sequences,
        or None if all sequences share the same first symbol (no split needed).
    """
    if not symbol_seqs:
        return None

    first_symbols = {}
    for seq in symbol_seqs:
        if seq:
            first_symbols.setdefault(seq[0], []).append(seq)
        else:
            first_symbols.setdefault('(empty)', []).append(seq)

    # Only split if 2+ sub-groups have enough methods
    viable = {k: v for k, v in first_symbols.items() if len(v) >= min_subgroup}

    if len(viable) <= 1:
        return None  # no useful split

    return viable


def _recursive_split(symbol_seqs, min_subgroup=3, max_depth=3, _depth=0):
    """Recursively split by first symbol until sub-groups are uniform.

    Returns:
        dict mapping "first1.first2..." → list of sequences (leaf groups).
    """
    if _depth >= max_depth:
        return {"": symbol_seqs}

    splits = _split_by_first_symbol(symbol_seqs, min_subgroup=min_subgroup)
    if splits is None:
        return {"": symbol_seqs}

    result = {}
    for first_sym, sub_seqs in splits.items():
        sub_leaves = _recursive_split(sub_seqs, min_subgroup, max_depth, _depth + 1)
        for suffix, leaf_seqs in sub_leaves.items():
            key = f"{first_sym}.{suffix}" if suffix else first_sym
            result[key] = leaf_seqs
    return result


def _infer_group(label, group_seqs, group_files, project_root, min_coverage, prefer, kmax, N, include_kore=False, include_idregex=False, method='langsize', min_methods=3, crx_method='standard', min_structure=0.0, split_mixed=False):
    """Infer grammar for one package group. Module-level for ProcessPoolExecutor."""
    filtered = frequency_filter(group_seqs, min_coverage=min_coverage)
    imports = _extract_imports(group_files)
    arg_patterns = _build_arg_patterns(group_files)
    packages = _top_packages(group_files, project_root)
    symbol_seqs = [[text for _, text, _ in seq] for seq in filtered]

    # Diversity threshold: skip if too few methods or too diverse
    n_methods = len(symbol_seqs)
    if n_methods < min_methods:
        meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages, "skip_reason": "too_few_methods"}
        return (label, None, len(filtered), meta)

    unique_seqs = len(set(tuple(s) for s in symbol_seqs))
    unique_ratio = unique_seqs / n_methods
    if unique_ratio > 0.9:
        meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages, "skip_reason": "too_diverse"}
        return (label, None, len(filtered), meta)

    # Split mixed-pattern groups before CRX
    if split_mixed:
        leaves = _recursive_split(symbol_seqs, min_subgroup=min_methods, max_depth=3)
        if len(leaves) > 1:
            # Infer each leaf, return ALL that pass
            all_results = []
            total_count = 0
            for leaf_key, leaf_seqs in leaves.items():
                leaf_label = f"{label} [{leaf_key}]" if leaf_key else label
                leaf_result = infer_ensemble(leaf_seqs, kmax=kmax, N=N, prefer=prefer, min_coverage=min_coverage, include_kore=include_kore, include_idregex=include_idregex, method=method)
                total_count += len(leaf_seqs)
                if leaf_result and leaf_result.get('best') and leaf_result['best'].get('grammar'):
                    g = leaf_result['best']['grammar']
                    ok, _ = validate_sore(g)
                    if ok:
                        if min_structure > 0 and grammar_structure_score(g) < min_structure:
                            continue
                        all_results.append((leaf_label, leaf_result, len(leaf_seqs)))

            if all_results:
                # Return the best one, but store all in meta for later use
                best_label, best_result, best_count = max(all_results, key=lambda x: grammar_structure_score(x[1]['best']['grammar']))
                meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages,
                        "split": True, "n_leaves": len(leaves), "n_grammars": len(all_results),
                        "all_grammars": [(l, r['best']['grammar'], grammar_structure_score(r['best']['grammar']), c) for l, r, c in all_results]}
                return (best_label, best_result, total_count, meta)
            # fall through to unsplit inference

    if crx_method == 'refined':
        from ..crx_refined import crx_with_confidence
        info = crx_with_confidence(symbol_seqs)
        result = {
            'best': {'algorithm': 'CRX-refined', 'grammar': info['grammar'], 'mdl_score': info['confidence']},
            'all': [{'algorithm': 'CRX-refined', 'grammar': info['grammar'], 'mdl_score': info['confidence']}],
            'why': f"CRX-refined (confidence={info['confidence']:.2f})",
        }
    else:
        result = infer_ensemble(symbol_seqs, kmax=kmax, N=N, prefer=prefer, min_coverage=min_coverage, include_kore=include_kore, include_idregex=include_idregex)

    # Validate grammar is parseable
    if result and result.get('best') and result['best'].get('grammar'):
        grammar = result['best']['grammar']
        ok, err = validate_sore(grammar)
        if not ok:
            meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages, "skip_reason": "malformed_grammar", "skip_detail": err}
            return (label, None, len(filtered), meta)
        if min_structure > 0 and grammar_structure_score(grammar) < min_structure:
            meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages, "skip_reason": "low_structure", "structure_score": grammar_structure_score(grammar)}
            return (label, None, len(filtered), meta)

    meta = {"files": group_files, "imports": imports, "arg_patterns": arg_patterns, "packages": packages}
    return (label, result, len(filtered), meta)


def analyze_by_package(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, min_pkg_size=3, include_kore=False, include_idregex=False, method='langsize', min_methods=3, crx_method='standard', min_structure=0.0, split_mixed=False):
    """Preprocess and group by package directory, infer per group.

    Groups methods by their file's relative directory path, merging
    small packages (< min_pkg_size methods) upward to their parent.

    Returns:
        list of (package_label, ensemble_result_dict, method_count, meta).
    """
    t0 = time.time()
    sequences, seq_files = _preprocess_files(file_paths)
    if not sequences:
        return []
    _vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")

    seq_packages = [_file_to_package(fp, project_root) for fp in seq_files]
    groups, ungrouped = _group_by_package(
        list(enumerate(seq_packages)), min_size=min_pkg_size
    )

    _vprint(f"Groups: {len(groups)} named, {len(ungrouped)} ungrouped methods")
    for label, idxs in groups:
        _vprint(f"  ├ {label} ({len(idxs)} methods)")
    if ungrouped:
        _vprint(f"  └ (other) ({len(ungrouped)} methods)")

    results = []
    n_workers = os.cpu_count()
    _vprint(f"Inferring {len(groups)} groups across {n_workers} workers ...")
    with ProcessPoolExecutor(max_workers=n_workers) as ex:
        futures = {}
        for label, indices in groups:
            gs = [sequences[i] for i in indices]
            gf = set(seq_files[i] for i in indices)
            f = ex.submit(_infer_group, label, gs, gf, project_root,
                           min_coverage, prefer, kmax, N, include_kore, include_idregex, method, min_methods, crx_method, min_structure, split_mixed)
            futures[f] = label

        for f in as_completed(futures):
            label, result, count, meta = f.result()
            results.append((label, result, count, meta))
            _vprint(f"Infer {label} ({count} methods) done")

    results.sort(key=lambda x: x[0])

    if ungrouped:
        ungrouped_files = set(seq_files[i] for i in ungrouped)
        ungrouped_seqs = [sequences[i] for i in ungrouped]
        results.append(("(other)", None, len(ungrouped_seqs), {"files": ungrouped_files, "imports": [], "arg_patterns": {}, "packages": []}))

    return results


def infer(file_paths, extension, min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False, method='langsize'):
    """Run full pipeline: preprocess → frequency filter → ensemble infer.

    Args:
        file_paths: list of source file paths (same language).
        extension: language extension (e.g. '.py').
        min_coverage: BEX core coverage threshold for outlier removal.
        prefer: inference algorithm preference ('crx', 'idregex', or None).
        kmax: max k for k-ORE algorithms.
        N: number of random trials.

    Returns:
        Ensemble result dict from infer_ensemble.
    """
    sequences, _ = _preprocess_files(file_paths)
    sequences = frequency_filter(sequences, min_coverage=min_coverage)

    symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]

    return infer_ensemble(symbol_seqs, kmax=kmax, N=N, prefer=prefer, min_coverage=min_coverage, method=method)


def _merge_up(pkg):
    """Go one directory level up from a package path."""
    parts = pkg.replace(os.sep, "/").rstrip("/").split("/")
    if len(parts) <= 1:
        return ""
    return "/".join(parts[:-1])


def _group_by_package(indices_and_packages, min_size=3):
    """Group method indices by package, merging small packages upward.

    Small packages that would merge into root are discarded (too small
    to form a meaningful group and not worth root-level inference).

    Args:
        indices_and_packages: list of (index, package_path) tuples.
        min_size: minimum methods to keep a standalone group.

    Returns:
        (groups, ungrouped) where groups is [(label, [indices])] sorted by label,
        and ungrouped is a list of indices that couldn't form a named group.
    """
    pkg_to_indices = {}
    for idx, pkg in indices_and_packages:
        pkg_to_indices.setdefault(pkg, []).append(idx)

    ungrouped = []

    while True:
        to_merge = {p for p, idxs in pkg_to_indices.items()
                    if p != "" and len(idxs) < min_size}
        if not to_merge:
            break
        new_groups = {}
        for pkg, indices in pkg_to_indices.items():
            if pkg in to_merge:
                parent = _merge_up(pkg)
                if parent == "":
                    ungrouped.extend(indices)
                else:
                    new_groups.setdefault(parent, []).extend(indices)
            else:
                new_groups.setdefault(pkg, []).extend(indices)
        pkg_to_indices = new_groups

    return sorted(pkg_to_indices.items(), key=lambda x: x[0]), ungrouped


def _filter_glob(files, include=None, exclude=None):
    """Filter file list by include/exclude glob patterns."""
    if include:
        spec = pathspec.PathSpec.from_lines("gitwildmatch", [include])
        files = [f for f in files if spec.match_file(f)]
    if exclude:
        spec = pathspec.PathSpec.from_lines("gitwildmatch", [exclude])
        files = [f for f in files if not spec.match_file(f)]
    return files


def analyze_by_reduce(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False, include_idregex=False, method='langsize', min_methods=3, crx_method='standard', min_structure=0.0, reduce_threshold=0.15):
    """Reduce-style analysis: group by directory, then merge similar groups.

    Uses Algorithm 4 (Reduce, TODS 2010) to merge directories with similar
    calling patterns. This finds natural groupings — directories that share
    the same calling convention get merged into larger groups.

    Args:
        reduce_threshold: similarity threshold for Reduce (0.05=conservative, 0.15=moderate, 0.30=aggressive)
    """
    from bex.reduce import reduce_and_infer

    t0 = time.time()
    sequences, seq_files = _preprocess_files(file_paths)
    if not sequences:
        return []
    _vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")

    # Build initial contexts by directory
    seq_packages = [_file_to_package(fp, project_root) for fp in seq_files]
    initial_contexts = {}
    for i, pkg in enumerate(seq_packages):
        initial_contexts.setdefault(pkg, []).append(sequences[i])

    _vprint(f"Initial contexts: {len(initial_contexts)} directories")

    # Run Reduce to merge similar contexts
    t1 = time.time()
    result = reduce_and_infer(initial_contexts, reduce_threshold, min_methods=min_methods)
    _vprint(f"Reduce: {result['merge_info']['contexts_before']} → {result['merge_info']['contexts_after']} contexts ({time.time()-t1:.1f}s)")
    _vprint(f"  Merges: {result['merge_info']['merges']}, Minimized: {result['minimize_info']['merged_by_sore']}")

    # Convert to standard pipeline format
    results = []
    n_workers = os.cpu_count()
    _vprint(f"Inferring {len(result['merged'])} contexts across {n_workers} workers ...")
    with ProcessPoolExecutor(max_workers=n_workers) as ex:
        futures = {}
        for label, seqs in result['merged'].items():
            f = ex.submit(_infer_group, label, seqs, set(), project_root,
                           min_coverage, prefer, kmax, N, include_kore, include_idregex, method, min_methods, crx_method, min_structure)
            futures[f] = label

        done = 0
        for f in as_completed(futures):
            done += 1
            if done % 20 == 0 or done == len(futures):
                _vprint(f"  [{done}/{len(futures)}]")
            results.append(f.result())

    return results


def analyze_by_ilocal(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False, include_idregex=False, method='langsize', min_methods=3, crx_method='standard', min_structure=0.0, context_strategy="dir", reduce=True):
    """iLocal-style analysis: extract (context, sequence) pairs, reduce, infer.

    Instead of hard-coding directory as grouping key, this extracts contexts
    using a configurable strategy, then uses iLocal's reduce to merge contexts
    with identical signatures.

    Args:
        context_strategy: one of 'dir', 'file', 'parent_dir', 'depth_2',
                         'depth_3', 'imports', 'symbol_overlap'
        reduce: whether to run iLocal reduce to merge similar contexts
    """
    from .ilocal_source import CONTEXT_STRATEGIES, reduce_contexts

    t0 = time.time()
    sequences, seq_files = _preprocess_files(file_paths)
    if not sequences:
        return []
    _vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")

    strategy_fn = CONTEXT_STRATEGIES.get(context_strategy)
    if not strategy_fn:
        _vprint(f"Unknown context strategy: {context_strategy}, falling back to 'dir'")
        strategy_fn = CONTEXT_STRATEGIES["dir"]

    t1 = time.time()
    context_groups = strategy_fn(sequences, seq_files, project_root)
    _vprint(f"Context extraction ({context_strategy}): {len(context_groups)} contexts ({time.time()-t1:.1f}s)")

    if reduce:
        t2 = time.time()
        before = len(context_groups)
        context_groups = reduce_contexts(context_groups)
        _vprint(f"Reduce: {before} → {len(context_groups)} contexts ({time.time()-t2:.1f}s)")

    _vprint(f"Contexts: {len(context_groups)} groups")
    for label, seqs in sorted(context_groups.items(), key=lambda x: -len(x[1]))[:10]:
        _vprint(f"  ├ {label[:60]} ({len(seqs)} methods)")
    if len(context_groups) > 10:
        _vprint(f"  └ ... and {len(context_groups) - 10} more")

    results = []
    n_workers = os.cpu_count()
    _vprint(f"Inferring {len(context_groups)} contexts across {n_workers} workers ...")
    with ProcessPoolExecutor(max_workers=n_workers) as ex:
        futures = {}
        for label, seqs in context_groups.items():
            f = ex.submit(_infer_group, label, seqs, set(), project_root,
                           min_coverage, prefer, kmax, N, include_kore, include_idregex, method, min_methods, crx_method, min_structure)
            futures[f] = label

        done = 0
        for f in as_completed(futures):
            done += 1
            if done % 20 == 0 or done == len(futures):
                _vprint(f"  [{done}/{len(futures)}]")
            results.append(f.result())

    return results


def analyze_directory(
    dir_path,
    min_coverage=DEFAULT_COVERAGE,
    prefer=None,
    kmax=2,
    slice="flat",
    include=None,
    exclude=None,
    main_only=False,
    include_kore=False,
    include_idregex=False,
    method='langsize',
    min_methods=3,
    crx_method='standard',
    min_structure=0.0,
    context_strategy="dir",
    reduce_threshold=0.15,
    split_mixed=False,
):
    """Scan a directory and run analysis for each language found.

    Args:
        dir_path: directory to scan.
        min_coverage: BEX core coverage threshold for outlier removal.
        prefer: algorithm preference.
        kmax: max k for k-ORE algorithms.
        slice: grouping strategy — "flat" (one per language) or "package" (per directory).
        include: optional glob — only process files matching this pattern.
        exclude: optional glob — skip files matching this pattern.
        main_only: exclude test files when True.

    Returns:
        dict mapping extension → list of (label, result_dict, count, meta) tuples.
    """
    groups = scan_directory(dir_path)
    results = {}
    for ext, files in groups.items():
        if len(files) < 1:
            continue
        files = _filter_glob(files, include=include, exclude=exclude)
        if main_only:
            files = [f for f in files if _is_main_source(f)]
        if not files:
            continue
        if slice == "package":
            results[ext] = analyze_by_package(
                files, ext,
                project_root=dir_path,
                min_coverage=min_coverage,
                prefer=prefer,
                kmax=kmax,
                include_kore=include_kore,
                include_idregex=include_idregex,
                method=method,
                min_methods=min_methods,
                crx_method=crx_method,
                min_structure=min_structure,
                split_mixed=split_mixed,
            )
        elif slice == "reduce":
            results[ext] = analyze_by_reduce(
                files, ext,
                project_root=dir_path,
                min_coverage=min_coverage,
                prefer=prefer,
                kmax=kmax,
                include_kore=include_kore,
                include_idregex=include_idregex,
                method=method,
                min_methods=min_methods,
                crx_method=crx_method,
                min_structure=min_structure,
                reduce_threshold=reduce_threshold,
            )
        elif slice == "ilocal":
            results[ext] = analyze_by_ilocal(
                files, ext,
                project_root=dir_path,
                min_coverage=min_coverage,
                prefer=prefer,
                kmax=kmax,
                include_kore=include_kore,
                include_idregex=include_idregex,
                method=method,
                min_methods=min_methods,
                crx_method=crx_method,
                min_structure=min_structure,
                context_strategy=context_strategy,
            )
        else:
            results[ext] = analyze_clusters(
                files, ext,
                project_root=dir_path,
                min_coverage=min_coverage,
                prefer=prefer,
                kmax=kmax,
                include_kore=include_kore,
                method=method,
                min_methods=min_methods,
            )
    return results


def _build_json_output(results):
    """Convert results dict to a compact JSON structure for prompt injection."""
    output = []
    for ext, clusters in results.items():
        lang = {"language": ext, "conventions": []}
        total_methods = 0
        for label, result, count, meta in clusters:
            total_methods += count
            entry = {
                "label": label,
                "method_count": count,
            }
            if result and result.get("best"):
                entry["algorithm"] = result["best"]["algorithm"]
                entry["grammar"] = result["best"]["grammar"]
                entry["mdl_score"] = round(result['best']['mdl_score'], 1)
            entry["imports"] = meta.get("imports", [])
            entry["arg_patterns"] = meta.get("arg_patterns", {})
            lang["conventions"].append(entry)
        lang["total_methods"] = total_methods
        output.append(lang)
    return json.dumps(output, indent=2)


def _build_yaml_output(results, dir_path, max_mdl=200.0):
    """Build YAML output grouped by top-level module, sorted by MDL.

    Filters out (other), no-grammar groups, and groups above max_mdl.
    Returns YAML string.
    """
    import yaml

    project_name = os.path.basename(os.path.abspath(dir_path))

    # Collect all entries grouped by top-level module
    modules = {}
    total_patterns = 0
    total_methods = 0

    for ext, clusters in results.items():
        for label, result, count, meta in clusters:
            total_methods += count
            if label == "(other)":
                continue
            if not result or not result.get("best"):
                continue
            best = result["best"]
            if best["mdl_score"] > max_mdl:
                continue

            # Extract top-level module from package path
            parts = label.replace(os.sep, "/").split("/")
            module = parts[0] if len(parts) > 1 else "(root)"

            # When recursive split produced leaf grammars, output each one
            all_grammars = meta.get("all_grammars", [])
            if all_grammars:
                for leaf_label, leaf_grammar, leaf_score, leaf_count in all_grammars:
                    entry = {
                        "package": leaf_label,
                        "methods": leaf_count,
                        "grammar": leaf_grammar,
                        "score": round(leaf_score, 3),
                        "algorithm": best["algorithm"],
                        "mdl": round(best["mdl_score"], 1),
                    }
                    modules.setdefault(module, []).append(entry)
                    total_patterns += 1
            else:
                entry = {
                    "package": label,
                    "methods": count,
                    "grammar": best["grammar"],
                    "score": round(best.get("mdl_score", 0), 3),
                    "algorithm": best["algorithm"],
                    "mdl": round(best["mdl_score"], 1),
                }

                modules.setdefault(module, []).append(entry)
                total_patterns += 1

    # Sort entries within each module by MDL
    for module in modules:
        modules[module].sort(key=lambda x: x["mdl"])

    # Sort modules by name
    ordered = dict(sorted(modules.items()))

    # Build YAML
    header = f"# {project_name} — {total_patterns} patterns ({total_methods} methods)\n\n"
    yaml_content = yaml.dump(ordered, default_flow_style=False, sort_keys=False, allow_unicode=True)
    return header + yaml_content


def _persist_grammars(yaml_content, dir_path):
    """Write grammars to {dir_path}/.dervish/grammars.yml"""
    dervish_dir = os.path.join(dir_path, ".dervish")
    os.makedirs(dervish_dir, exist_ok=True)
    out_path = os.path.join(dervish_dir, "grammars.yml")
    with open(out_path, "w") as f:
        f.write(yaml_content)
    _vprint(f"Persisted to {out_path}")


def _parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description="Analyze a directory of source code for behavioral conventions.",
    )
    parser.add_argument("directory", help="Directory to scan")
    parser.add_argument(
        "--prefer",
        choices=["crx", "idregex"],
        help="Skip ensemble, use only this algorithm",
    )
    parser.add_argument(
        "--kore", action="store_true",
        help="Include kORE in ensemble (off by default for speed)",
    )
    parser.add_argument(
        "--idregex", action="store_true",
        help="Include iDRegEx in ensemble (off by default — slow on large groups)",
    )
    parser.add_argument(
        "--kmax", type=int, default=2,
        help="Maximum k for k-ORE algorithms (default: 2)",
    )
    parser.add_argument(
        "--min-coverage", type=float, default=DEFAULT_COVERAGE,
        help="BEX core coverage threshold — outlier methods are removed until this fraction remains (default: 0.8)",
    )
    parser.add_argument(
        "--slice", choices=["flat", "package", "reduce", "ilocal"], default="flat",
        help="Grouping strategy: flat (one per language), package (per directory), reduce (merge similar dirs), or ilocal (iLocal context-based) (default: flat)",
    )
    parser.add_argument(
        "--reduce-threshold", type=float, default=0.15,
        help="Similarity threshold for reduce slicing (0.05=conservative, 0.15=moderate, 0.30=aggressive) (default: 0.15)",
    )
    parser.add_argument(
        "--context-strategy", choices=["dir", "file", "parent_dir", "depth_2", "depth_3", "imports", "symbol_overlap"], default="dir",
        help="Context extraction strategy for ilocal slicing (default: dir)",
    )
    parser.add_argument(
        "--include",
        help="Glob pattern to filter files (e.g. '**/src/test/**')",
    )
    parser.add_argument(
        "--exclude",
        help="Glob pattern to skip files (e.g. '**/build/**')",
    )
    parser.add_argument(
        "--format", choices=["text", "json"], default="text",
        help="Output format (default: text)",
    )
    parser.add_argument(
        "--json", action="store_true", dest="json_flag",
        help="Shortcut for --format json",
    )
    parser.add_argument(
        "--verbose", action="store_true",
        help="Print progress to stderr",
    )
    parser.add_argument(
        "--main-only", action="store_true",
        help="Exclude test files (src/test/**, *Test.*, etc.)",
    )
    parser.add_argument(
        "--scoring-method", choices=["langsize", "mdl"], default="langsize",
        help="Scoring method: langsize (default, Bex et al.) or mdl (fallback)",
    )
    parser.add_argument(
        "--min-methods", type=int, default=3,
        help="Minimum methods per group to infer grammar (default: 3). Groups with fewer are skipped.",
    )
    parser.add_argument(
        "--crx-method", choices=["standard", "refined"], default="standard",
        help="CRX method: standard (default) or refined (cluster-then-infer, tighter grammars)",
    )
    parser.add_argument(
        "--min-structure", type=float, default=0.0,
        help="Minimum structure score (0.0-1.0) to keep grammar. Flat bags of symbols below this are dropped (default: 0, keep all)",
    )
    parser.add_argument(
        "--split-mixed", action="store_true",
        help="Split groups with mixed first symbols before CRX inference (produces tighter grammars)",
    )
    return parser.parse_args(argv)


def main():
    args = _parse_args()
    global VERBOSE, _vstart
    VERBOSE = args.verbose
    _vstart = time.time()
    _vprint(f"Scanning {args.directory} ...")
    results = analyze_directory(
        args.directory,
        min_coverage=args.min_coverage,
        prefer=args.prefer,
        kmax=args.kmax,
        slice=args.slice,
        include=args.include,
        exclude=args.exclude,
        main_only=args.main_only,
        include_kore=args.kore,
        include_idregex=args.idregex,
        method=args.scoring_method,
        min_methods=args.min_methods,
        crx_method=args.crx_method,
        min_structure=args.min_structure,
        context_strategy=args.context_strategy,
        reduce_threshold=args.reduce_threshold,
        split_mixed=args.split_mixed,
    )

    if args.json_flag or args.format == "json":
        print(_build_json_output(results))
        return

    for ext, clusters in results.items():
        print(f"\n{ext}:")
        for label, result, count, meta in clusters:
            if result and result.get("best"):
                best = result["best"]
                print(f"  ╰─ {label} ({count} methods)")
                print(f"     Algorithm: {best['algorithm']}")
                print(f"     Grammar:   {best['grammar']}")
                print(f"     Score:     {best['mdl_score']}")
            else:
                reason = meta.get("skip_reason", "")
                suffix = f" — {reason}" if reason else ""
                print(f"  ╰─ {label} ({count} methods) — no grammar{suffix}")
            imps = meta.get("imports", [])
            if imps:
                joined = "  |  ".join(imps[:6])
                print(f"     Imports:  {joined}")
                if len(imps) > 6:
                    print(f"               ... and {len(imps) - 6} more")
            argp = meta.get("arg_patterns", {})
            if argp:
                top_calls = sorted(argp.items(), key=lambda x: -x[1]["occurrences"])[:4]
                for call_name, asum in top_calls:
                    ac = asum["arg_count"]
                    pats = asum["patterns"][:2]
                    pat_strs = [f"{p['args']}:{','.join(p['types'])}" for p in pats]
                    print(f"     Args({call_name}): n={ac['common']} "
                          f"[{'; '.join(pat_strs)}]")


if __name__ == "__main__":
    main()
