"""Dervish — MCP server.

Provides tools to infer regular expression grammars from example sequences,
and to look up the right grammar at code generation time.

Run as: python -m bex.mcp_server
"""

import os

from mcp.server.fastmcp import FastMCP

from .ensemble import infer_ensemble, _matches
from .grammar_index import load_grammar_index
from .tag_preprocessor.analyze import (
    analyze_directory as _analyze_directory,
    _build_yaml_output,
    _persist_grammars,
)

mcp = FastMCP("grammar-inference", log_level="ERROR")


# ---------------------------------------------------------------------------
# Inference tools
# ---------------------------------------------------------------------------

@mcp.tool()
def infer_best_grammar(
    sequences: list[list[str]],
    prefer: str = "",
    kmax: int = 2,
    N: int = 3,
    min_coverage: float = 1.0,
    method: str = "langsize",
) -> str:
    """Infer a compact grammar from example sequences. Use this when you
    have examples of sequential data and want to learn the pattern.

    The grammar compresses N examples into ~100 chars — far fewer tokens
    than passing all examples. Pass the existing sequences, get back a
    pattern you can follow to generate new instances.

    Runs CRX + iDRegEx, picks best by scoring. kORE is excluded by
    default (slow, rarely wins on real data). Set prefer='koreinference'
    to force it.

    Args:
        sequences: List of sequences, each a list of strings (symbols in
            the order they appear). Example: [["file","copy","command"],
            ["file","template","command"]].
        prefer: Optional — 'crx' for full vocabulary (accepts all examples),
            'idregex' for deterministic minimal core, 'koreinference' for
            k-OA with rwr0 repair (slow). Omit to auto-pick by MDL.
        kmax: Context depth for k-ORE inference (iDRegEx, kOREInference).
            Default 2.
        N: Random trials for k-ORE inference (higher = better, slower).
        min_coverage: (Expert) When < 1.0, also runs a **core+outlier analysis**:
            iteratively removes outlier sequences (those with rarest symbols)
            until at least this fraction remain. Returns the core grammar
            for the majority, plus a list of which sequences were removed and why.
            Default 1.0 = no core analysis. Set to 0.8 to find the tight
            pattern shared by ~80% of examples while flagging the other ~20%
            as variations.

    Returns:
        A formatted string with the best grammar, scores, and explanation.
        When min_coverage < 1.0, includes the core grammar and outlier info.
        Grammar notation: a.b = a then b, (a+b) = a or b, r? = optional,
        r+ = one or more, r+? = zero or more.
    """
    pref = prefer if prefer else None
    result = infer_ensemble(sequences, kmax=kmax, N=N, prefer=pref, min_coverage=min_coverage, method=method)
    if result['best'] is None:
        return f"No grammar found. {result['why']}"
    lines = [f"Best: {result['best']['algorithm']} (Score {result['best']['mdl_score']})",
             f"Grammar: {result['best']['grammar']}",
             ""]
    if len(result['all']) > 1:
        for r in result['all']:
            m = sum(1 for s in sequences if _matches(r['grammar'], s))
            lines.append(f"  {r['algorithm']:10s}  Score={r['mdl_score']:>8.2f}  match={m}/{len(sequences)}")
    lines.append("")
    lines.append(f"Why: {result['why']}")
    if 'core' in result and result['core']:
        c = result['core']
        lines.append(f"\nCore CRX ({c['coverage']:.0%} coverage, {c['outlier_count']} outliers): {c['grammar']}")
        if c['outliers']:
            lines.append(f"  Outlier sequences:")
            for i, o in enumerate(c['outliers'], 1):
                lines.append(f"    {i}. {' → '.join(str(x) for x in o[:8])}{'...' if len(o) > 8 else ''}")
    return "\n".join(lines)


@mcp.tool()
def analyze_directory(
    directory: str,
    slice: str = "package",
    min_coverage: float = 0.05,
    prefer: str = "",
    kmax: int = 2,
    include: str = "",
    exclude: str = "",
    main_only: bool = False,
    max_mdl: float = 200,
    persist: bool = True,
    method: str = "langsize",
    min_methods: int = 3,
    min_structure: float = 0.2,
    split_mixed: bool = True,
) -> str:
    """Scan a source code directory and infer behavioral conventions
    (regular expression grammars) per package. Returns compact patterns
    grouped by module, sorted by quality.

    Use this when you need to understand the calling conventions in a
    codebase — what patterns new code should follow. The grammar
    compresses each package's method call patterns into a compact
    regular expression.

    Auto-persists results to {directory}/.dervish/grammars.yml unless
    persist=False. An agent can later look up the right grammar via
    get_grammar() during code generation.

    Args:
        directory: Path to the source code directory to analyze.
        slice: Grouping strategy — 'package' (per directory, default)
            or 'flat' (one per language).
        min_coverage: Coverage threshold for outlier removal (0.01–1.0).
            Lower = see more symbols. Default 0.05.
        prefer: Optional — 'crx' for full vocabulary, 'idregex' for
            minimal core. Omit to auto-pick by scoring.
        kmax: Context depth for k-ORE inference. Default 2.
        include: Glob pattern to include only matching files.
        exclude: Glob pattern to skip matching files.
        main_only: When True, exclude test files. Default False.
        max_mdl: Drop groups with score above this threshold. Default 200.
        persist: When True (default), write results to
            {directory}/.dervish/grammars.yml.
        method: Scoring method — 'langsize' (default) or 'mdl'.
        min_methods: Minimum methods per group to attempt inference. Default 3.
        min_structure: Minimum grammar structure score (0.0–1.0). Groups
            below this produce flat bags. Default 0.2.
        split_mixed: When True (default), recursively split groups with
            diverse first symbols into uniform sub-groups before inference.

    Returns:
        YAML string with grammars grouped by top-level module, sorted
        by score (tightest/most useful first).
    """
    results = _analyze_directory(
        directory,
        min_coverage=min_coverage,
        prefer=prefer or None,
        kmax=kmax,
        slice=slice,
        include=include or None,
        exclude=exclude or None,
        main_only=main_only,
        method=method,
        min_methods=min_methods,
        split_mixed=split_mixed,
        min_structure=min_structure,
    )
    yaml_content = _build_yaml_output(results, directory, max_mdl=max_mdl)
    if persist:
        _persist_grammars(yaml_content, directory)
    return yaml_content


# ---------------------------------------------------------------------------
# Runtime lookup tools
# ---------------------------------------------------------------------------

@mcp.tool()
def get_grammar(
    directory: str,
    file_path: str,
    context_symbol: str = "",
) -> str:
    """Get the grammar constraint for generating code in a specific file.

    Call this before generating code to get the GBNF grammar that matches
    the file's package and the calling context. The agent should pass the
    returned grammar to the LLM's constrained generation backend.

    Workflow:
      1. Agent is about to write code in `file_path`
      2. Call get_grammar(directory, file_path, context_symbol="return")
      3. Pass the returned GBNF to the LLM's grammar-constrained sampler
      4. LLM generates code that follows the package's convention

    Args:
        directory: Project root (must match the directory used in
            analyze_directory).
        file_path: Path to the file being written (absolute or relative
            to directory).
        context_symbol: Optional — the first symbol of the code being
            generated (e.g. "return", "if", "try"). When provided, returns
            the leaf grammar for that specific context. When empty, returns
            the best grammar for the file's package.

    Returns:
        GBNF grammar string, or a message explaining why no grammar was found.
    """
    idx = load_grammar_index(directory)
    if not idx.get_all():
        return f"No grammars indexed for {directory}. Run analyze_directory first."

    ctx = context_symbol if context_symbol else None
    grammar = idx.get(file_path, context_symbol=ctx)
    if grammar is None:
        return f"No grammar found for {file_path} (context={context_symbol or 'best'})."

    return grammar


@mcp.tool()
def get_package_grammars(
    directory: str,
    file_path: str,
) -> str:
    """Get all available grammars for a file's package, ranked by quality.

    Use this to explore what conventions exist for a package before
    choosing which context to generate in. Returns all leaf grammars
    (one per calling context like "return", "if", "try", etc.).

    Args:
        directory: Project root.
        file_path: Path to the file (absolute or relative to directory).

    Returns:
        Formatted list of (context_symbol, grammar, score, methods) tuples.
    """
    idx = load_grammar_index(directory)
    if not idx.get_all():
        return f"No grammars indexed for {directory}. Run analyze_directory first."

    entries = idx.get_package(file_path)
    if not entries:
        return f"No grammars found for package of {file_path}."

    lines = [f"Grammars for {file_path}:"]
    for sym, grammar, score, methods in entries:
        label = f"[{sym}]" if sym else "(best)"
        lines.append(f"  {label:25s}  score={score:.3f}  {methods}m  {grammar[:70]}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    mcp.run()


if __name__ == "__main__":
    main()
