"""
Context Selection Engine — Core of Kodara (Phase 3).

Given a natural language query, the engine:
  1. Scores all modules via SemanticScorer (stemming + IDF + optional cosine)
  2. Selects the highest-scoring modules
  3. Expands selection via 1-hop dependency traversal
  4. Assembles a minimal structured context object

Design goal: Reduce context from ~900k tokens to ~20k tokens while
preserving architectural understanding of the relevant subsystem.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from ..memory.store import MemoryStore, KodaraIndex
from ..search.semantic import SemanticScorer


@dataclass
class ContextResult:
    """Structured context assembled by the engine."""
    query: str
    modules: list[dict]            # Selected module metadata
    summaries: list[str]           # Human-readable summary lines
    relevant_files: list[str]      # Ordered list of relevant file paths
    graph_fragment: dict           # Minimal subgraph of dependencies
    token_estimate: int = 0        # Rough token count of this context
    total_modules_searched: int = 0


class ContextEngine:
    """
    Architecture-aware context selection engine.

    Works against a loaded KodaraIndex — no file system access required
    at query time. All data comes from the pre-built memory store.
    """

    def __init__(self, store: MemoryStore):
        self.store = store

    def query(
        self,
        query: str,
        max_modules: int = 8,
        depth: int = 1,
    ) -> ContextResult:
        """
        Execute a context query against the indexed repository.

        Args:
            query: Natural language query (e.g. "How does auth work?")
            max_modules: Maximum modules to include in context
            depth: Dependency traversal depth for expansion

        Returns:
            ContextResult with assembled minimal context
        """
        index = self.store.load()
        if not index:
            raise ValueError(
                "Project not initialised. Run: kodara init"
            )

        # 1. Score every module against the query (semantic: stem + IDF + cosine)
        scorer = SemanticScorer(index.modules)
        scored = scorer.score(query)

        # 2. Pick the top seed modules
        seeds_count = max(1, max_modules // 2)
        seed_paths = [path for _, path in scored[:seeds_count]]

        # 3. Expand seeds via dependency graph
        expanded = self._expand_via_graph(index, seed_paths, depth, max_modules)

        # 4. Assemble context
        return self._assemble_context(index, query, expanded, scored)

    # ── Expansion ────────────────────────────────────────────────────────────

    def _expand_via_graph(
        self,
        index: KodaraIndex,
        seeds: list[str],
        depth: int,
        max_total: int,
    ) -> list[str]:
        """
        BFS from seed modules collecting related modules within `depth` hops.
        Prioritizes seeds, then their dependencies, then their dependents.
        """
        selected: list[str] = list(seeds)
        visited: set[str] = set(seeds)

        queue = [(path, 0) for path in seeds]

        while queue and len(selected) < max_total:
            node, d = queue.pop(0)
            if d >= depth:
                continue

            mod = index.modules.get(node, {})

            # Visit dependencies first (what this module uses)
            for dep in mod.get("dependencies", []):
                if dep not in visited and dep in index.modules:
                    visited.add(dep)
                    selected.append(dep)
                    queue.append((dep, d + 1))
                    if len(selected) >= max_total:
                        break

            # Then dependents (what uses this module)
            for dep in mod.get("dependents", []):
                if dep not in visited and dep in index.modules:
                    visited.add(dep)
                    selected.append(dep)
                    queue.append((dep, d + 1))
                    if len(selected) >= max_total:
                        break

        return selected[:max_total]

    # ── Assembly ─────────────────────────────────────────────────────────────

    def _assemble_context(
        self,
        index: KodaraIndex,
        query: str,
        paths: list[str],
        scored: list[tuple[float, str]],
    ) -> ContextResult:
        """Assemble the final ContextResult from selected paths."""
        paths_set = set(paths)
        modules_out: list[dict] = []
        summaries: list[str] = []

        for path in paths:
            mod = index.modules.get(path, {})
            module_name = mod.get("module_name", path)
            summary = mod.get("summary", "")

            modules_out.append({
                "file": path,
                "module": module_name,
                "language": mod.get("language", ""),
                "classes": mod.get("classes", []),
                "functions": mod.get("functions", [])[:12],
                "exports": mod.get("exports", [])[:8],
                "imports": mod.get("imports", [])[:10],
                "dependencies": [d for d in mod.get("dependencies", []) if d in paths_set],
                "summary": summary,
            })

            if summary:
                summaries.append(f"**{module_name}**: {summary}")

        # Build minimal graph fragment
        graph_fragment = {
            "nodes": [
                {"id": p, "module": index.modules.get(p, {}).get("module_name", p)}
                for p in paths
            ],
            "edges": [
                e for e in index.graph_edges
                if e.get("from") in paths_set and e.get("to") in paths_set
            ],
        }

        # Rough token estimate: ~4 chars per token
        context_text = str(modules_out) + str(summaries)
        token_estimate = len(context_text) // 4

        return ContextResult(
            query=query,
            modules=modules_out,
            summaries=summaries,
            relevant_files=paths,
            graph_fragment=graph_fragment,
            token_estimate=token_estimate,
            total_modules_searched=len(index.modules),
        )

    # ── Formatting ───────────────────────────────────────────────────────────

    def format_context(self, result: ContextResult) -> str:
        """Format a ContextResult as readable markdown for LLM consumption."""
        reduction = "N/A"
        if result.total_modules_searched > 0:
            pct = 100 * (1 - len(result.modules) / result.total_modules_searched)
            reduction = f"{pct:.0f}%"

        lines = [
            f"# Kodara Context",
            f"**Query:** {result.query}",
            f"**Modules selected:** {len(result.modules)} / {result.total_modules_searched} "
            f"(context reduction: {reduction})",
            f"**Estimated tokens:** ~{result.token_estimate:,}",
            "",
        ]

        if result.summaries:
            lines.append("## Module Summaries")
            for s in result.summaries:
                lines.append(f"- {s}")
            lines.append("")

        lines.append("## Architecture")
        for mod in result.modules:
            lines.append(f"\n### `{mod['file']}`")
            lines.append(f"**Module:** `{mod['module']}` · **Language:** {mod['language']}")
            if mod["classes"]:
                lines.append(f"**Classes:** {', '.join(mod['classes'])}")
            if mod["functions"]:
                lines.append(f"**Functions:** {', '.join(mod['functions'][:8])}")
            if mod["exports"]:
                lines.append(f"**Exports:** {', '.join(mod['exports'][:6])}")
            if mod["imports"]:
                lines.append(f"**Imports:** {', '.join(mod['imports'][:8])}")
            if mod["dependencies"]:
                deps_str = ", ".join(mod["dependencies"])
                lines.append(f"**Internal deps:** {deps_str}")

        if result.graph_fragment["edges"]:
            lines.append("\n## Dependency Flow")
            for edge in result.graph_fragment["edges"]:
                lines.append(f"- `{edge['from']}` -> `{edge['to']}`")

        return "\n".join(lines)
