"""Grammar index for runtime lookup.

Loads grammars.yml and builds a lookup index mapping
(package, context_symbol) → GBNF grammar string.

Usage:
    from bex.grammar_index import load_grammar_index

    idx = load_grammar_index("/path/to/project")
    grammar = idx.get(file_path="src/services/Chat.kt", context_symbol="return")
    all_grammars = idx.get_package("src/services")
"""

import os
import re
import yaml


_LABEL_RE = re.compile(r"^(.+)\s+\[([^\]]+)\]$")


class GrammarIndex:
    """In-memory index of grammars keyed by (package, context_symbol)."""

    def __init__(self, project_root, entries):
        """
        Args:
            project_root: absolute path to the project root.
            entries: list of dicts with keys: package, grammar, score, methods, algorithm.
        """
        self.project_root = project_root
        self._by_package = {}       # package → [(symbol, grammar, score)]
        self._all = entries

        for e in entries:
            label = e["package"]
            m = _LABEL_RE.match(label)
            if m:
                pkg = m.group(1).rstrip("/")
                symbol = m.group(2)
            else:
                pkg = label.rstrip("/")
                symbol = ""

            self._by_package.setdefault(pkg, []).append((
                symbol,
                e["grammar"],
                e.get("score", 0),
                e.get("methods", 0),
            ))

        # Sort each package's entries by score descending (best first)
        for pkg in self._by_package:
            self._by_package[pkg].sort(key=lambda x: -x[2])

    def get(self, file_path, context_symbol=None):
        """Get the best grammar for a file, optionally filtered by context.

        Args:
            file_path: path to the source file (absolute or relative to project_root).
            context_symbol: if provided, match the leaf grammar whose first symbol
                is this. If None, return the best grammar for the package.

        Returns:
            GBNF grammar string, or None if no match.
        """
        pkg = self._resolve_package(file_path)
        entries = self._by_package.get(pkg, [])
        if not entries:
            return None

        if context_symbol:
            for sym, grammar, score, methods in entries:
                if sym == context_symbol:
                    return grammar

        # Fall back to best grammar for the package
        return entries[0][1] if entries else None

    def get_package(self, file_path):
        """Get all grammars for a file's package.

        Returns:
            list of (context_symbol, grammar, score, methods) tuples,
            sorted by score descending. Empty list if no match.
        """
        pkg = self._resolve_package(file_path)
        return list(self._by_package.get(pkg, []))

    def get_all(self):
        """Return all entries as a flat list."""
        return list(self._all)

    def packages(self):
        """Return sorted list of all indexed packages."""
        return sorted(self._by_package.keys())

    def _resolve_package(self, file_path):
        """Map a file path to its package."""
        # Make absolute if relative
        if not os.path.isabs(file_path):
            file_path = os.path.join(self.project_root, file_path)

        # Get directory of file relative to project root
        try:
            rel = os.path.relpath(os.path.dirname(file_path), self.project_root)
        except ValueError:
            # Different drives on Windows
            return ""
        if rel == ".":
            return ""
        return rel

    def __repr__(self):
        return f"GrammarIndex({self.project_root}, {len(self._all)} entries, {len(self._by_package)} packages)"


def load_grammar_index(project_root):
    """Load grammars.yml from {project_root}/.dervish/grammars.yml.

    Returns GrammarIndex, or empty index if no file found.
    """
    yml_path = os.path.join(project_root, ".dervish", "grammars.yml")
    if not os.path.exists(yml_path):
        return GrammarIndex(project_root, [])

    with open(yml_path) as f:
        data = yaml.safe_load(f)

    entries = []
    if isinstance(data, dict):
        for module, items in data.items():
            if not isinstance(items, list):
                continue
            for item in items:
                if isinstance(item, dict) and "package" in item and "grammar" in item:
                    entries.append(item)

    return GrammarIndex(project_root, entries)