"""Code queue context enrichment -- discover sibling dependencies,
channel discussion context, and prior art before task execution.

Ported from BOSS smack_code_queue_enrichment.  Adapted for Cohort's
flat-file, zero-infrastructure design.

Enrichment layers (each is best-effort, failures produce empty results):

1. **Sibling contracts** -- for NEW files: scan siblings for imports of
   the new module to discover what callers expect.
2. **Dependent files** -- for MODIFY tasks: find who imports the file
   being changed (backward-compat checking).
3. **Channel discussion** -- summarise the source channel conversation
   that led to this task (~500 token budget).
4. **File window** -- for large files: extract focused reading windows
   using simple heuristics (header + target section).

Usage::

    from cohort.pipeline.enrichment import TaskEnricher

    enricher = TaskEnricher(project_root=Path("."))

    context = enricher.enrich(
        file_manifest=["src/routes/users.py"],
        task_description="Add GET /users endpoint",
        channel_messages=[...],   # optional
    )
    # => EnrichedContext(sibling_contracts=[...], dependents=[...], ...)
    # => context.to_prompt_block()  # ~2000 tokens for injection into brief
"""

from __future__ import annotations

import ast
import logging
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Token budgets (approximate: 1 token ≈ 4 chars)
# ---------------------------------------------------------------------------

CODE_ENRICHMENT_CHAR_BUDGET = 8000     # ~2000 tokens
DISCUSSION_CHAR_BUDGET = 2000          # ~500 tokens
MAX_SIBLINGS = 10
MAX_DEPENDENTS = 10
MAX_USAGE_LINES = 10
LARGE_FILE_THRESHOLD = 200            # lines


# ---------------------------------------------------------------------------
# Output types
# ---------------------------------------------------------------------------

@dataclass
class SiblingContract:
    """What a sibling file expects to import from our target."""

    sibling_path: str
    expected_imports: List[str] = field(default_factory=list)
    usage_lines: List[str] = field(default_factory=list)


@dataclass
class DependentFile:
    """A file that imports from our target (backward-compat concern)."""

    path: str
    import_lines: List[str] = field(default_factory=list)


@dataclass
class FileWindow:
    """Focused reading window for a large file."""

    path: str
    header: str = ""       # imports / docstring section
    target: str = ""       # section around a target function/class
    total_lines: int = 0


@dataclass
class EnrichedContext:
    """All enrichment results for a task."""

    sibling_contracts: List[SiblingContract] = field(default_factory=list)
    dependents: List[DependentFile] = field(default_factory=list)
    discussion_summary: str = ""
    file_windows: List[FileWindow] = field(default_factory=list)

    def to_prompt_block(self) -> str:
        """Render enrichment as a prompt section within token budget."""
        sections: List[str] = []

        if self.sibling_contracts:
            lines = ["## Sibling Expectations"]
            for sc in self.sibling_contracts:
                lines.append(f"**{sc.sibling_path}** expects: {', '.join(sc.expected_imports)}")
                for ul in sc.usage_lines[:3]:
                    lines.append(f"  > {ul.strip()}")
            sections.append("\n".join(lines))

        if self.dependents:
            lines = ["## Files That Depend On Your Target"]
            for dep in self.dependents:
                lines.append(f"- **{dep.path}**: {'; '.join(dep.import_lines[:3])}")
            sections.append("\n".join(lines))

        if self.discussion_summary:
            sections.append(f"## Discussion Context\n{self.discussion_summary}")

        if self.file_windows:
            for fw in self.file_windows:
                lines = [f"## File Preview: {fw.path} ({fw.total_lines} lines)"]
                if fw.header:
                    lines.append("### Header\n```\n" + fw.header + "\n```")
                if fw.target:
                    lines.append("### Target Section\n```\n" + fw.target + "\n```")
                sections.append("\n".join(lines))

        result = "\n\n".join(sections)
        # Enforce budget
        if len(result) > CODE_ENRICHMENT_CHAR_BUDGET + DISCUSSION_CHAR_BUDGET:
            result = result[: CODE_ENRICHMENT_CHAR_BUDGET + DISCUSSION_CHAR_BUDGET]
            result += "\n[...truncated to fit token budget]"
        return result


# ---------------------------------------------------------------------------
# TaskEnricher
# ---------------------------------------------------------------------------

class TaskEnricher:
    """Discover contextual enrichment for pipeline tasks.

    Parameters
    ----------
    project_root:
        Root directory of the project being worked on.
    condenser:
        Optional ``ChannelCondenser`` instance for discussion summarisation.
        If None, discussion context uses raw truncation.
    """

    def __init__(
        self,
        project_root: Path,
        condenser: Any = None,
    ) -> None:
        self.root = project_root.resolve()
        self._condenser = condenser

    def enrich(
        self,
        file_manifest: List[str],
        task_description: str = "",
        channel_messages: Optional[Sequence[Dict[str, Any]]] = None,
        agent_id: str = "",
    ) -> EnrichedContext:
        """Build enriched context for a task.

        Every layer is best-effort -- failures produce empty results,
        never exceptions.
        """
        ctx = EnrichedContext()

        for filepath in file_manifest:
            resolved = self._resolve(filepath)
            if resolved is None:
                continue

            if resolved.exists():
                # Existing file -> find dependents + file window
                ctx.dependents.extend(self._find_dependents(resolved))
                window = self._build_file_window(resolved)
                if window:
                    ctx.file_windows.append(window)
            else:
                # New file -> find sibling contracts
                ctx.sibling_contracts.extend(self._find_sibling_contracts(resolved))

        # Channel discussion summary
        if channel_messages:
            ctx.discussion_summary = self._summarise_discussion(
                channel_messages, agent_id,
            )

        return ctx

    # ------------------------------------------------------------------
    # Layer 1: Sibling contracts (new files)
    # ------------------------------------------------------------------

    def _find_sibling_contracts(self, target: Path) -> List[SiblingContract]:
        """Scan siblings for imports of the not-yet-existing target module."""
        results: List[SiblingContract] = []
        if not target.parent.exists():
            return results

        module_name = target.stem  # e.g. "users" from "users.py"

        siblings = sorted(target.parent.glob("*.py"))[:MAX_SIBLINGS]
        for sib in siblings:
            if sib == target or sib.name.startswith("__"):
                continue
            try:
                contract = self._extract_imports_of(sib, module_name)
                if contract:
                    results.append(contract)
            except Exception:
                continue

        return results

    def _extract_imports_of(self, source: Path, module_name: str) -> Optional[SiblingContract]:
        """Check if *source* imports *module_name* and extract usage."""
        try:
            text = source.read_text(encoding="utf-8", errors="replace")
        except Exception:
            return None

        # Quick check before parsing
        if module_name not in text:
            return None

        expected: List[str] = []
        usage_lines: List[str] = []

        # Parse imports via AST
        try:
            tree = ast.parse(text, filename=str(source))
            for node in ast.walk(tree):
                if isinstance(node, ast.ImportFrom):
                    if node.module and module_name in node.module:
                        for alias in node.names:
                            expected.append(alias.name)
                elif isinstance(node, ast.Import):
                    for alias in node.names:
                        if module_name in alias.name:
                            expected.append(alias.name)
        except SyntaxError:
            # Fall back to regex
            for line in text.splitlines():
                m = re.match(rf"from\s+\S*{re.escape(module_name)}\S*\s+import\s+(.+)", line)
                if m:
                    expected.extend(name.strip() for name in m.group(1).split(","))

        if not expected:
            return None

        # Find usage lines
        for line in text.splitlines():
            for name in expected:
                if name in line and not line.strip().startswith(("import ", "from ")):
                    usage_lines.append(line)
                    if len(usage_lines) >= MAX_USAGE_LINES:
                        break
            if len(usage_lines) >= MAX_USAGE_LINES:
                break

        rel_path = (
            str(source.relative_to(self.root))
            if source.is_relative_to(self.root)
            else str(source)
        )
        return SiblingContract(
            sibling_path=rel_path,
            expected_imports=expected,
            usage_lines=usage_lines,
        )

    # ------------------------------------------------------------------
    # Layer 2: Dependent files (existing files)
    # ------------------------------------------------------------------

    def _find_dependents(self, target: Path) -> List[DependentFile]:
        """Find Python files that import from *target*."""
        results: List[DependentFile] = []
        module_name = target.stem

        # Search in the target's directory and one level up
        search_dirs = [target.parent]
        if target.parent.parent != self.root.parent:
            search_dirs.append(target.parent.parent)

        seen: set = set()
        for search_dir in search_dirs:
            if not search_dir.exists():
                continue
            for py_file in sorted(search_dir.rglob("*.py")):
                if py_file == target or str(py_file) in seen:
                    continue
                seen.add(str(py_file))
                try:
                    text = py_file.read_text(encoding="utf-8", errors="replace")
                except Exception:
                    continue
                if module_name not in text:
                    continue

                import_lines: List[str] = []
                for line in text.splitlines():
                    stripped = line.strip()
                    if (stripped.startswith(("import ", "from "))
                            and module_name in stripped):
                        import_lines.append(stripped)

                if import_lines:
                    rel = (
                        str(py_file.relative_to(self.root))
                        if py_file.is_relative_to(self.root)
                        else str(py_file)
                    )
                    results.append(DependentFile(path=rel, import_lines=import_lines))

                if len(results) >= MAX_DEPENDENTS:
                    return results

        return results

    # ------------------------------------------------------------------
    # Layer 3: Channel discussion summary
    # ------------------------------------------------------------------

    def _summarise_discussion(
        self,
        messages: Sequence[Dict[str, Any]],
        agent_id: str,
    ) -> str:
        """Summarise channel discussion that led to this task."""
        # Filter out noise (system messages, status updates)
        useful = [
            m for m in messages
            if m.get("sender", "") not in ("system", "worker")
            and len(m.get("content", "")) > 10
        ]

        if not useful:
            return ""

        # Use condenser if available
        if self._condenser:
            try:
                ctx = self._condenser.condense(useful[-20:], agent_id=agent_id)
                block = ctx.to_prompt_block()
                return block[:DISCUSSION_CHAR_BUDGET]
            except Exception:
                pass

        # Raw truncation fallback
        lines: List[str] = []
        budget = DISCUSSION_CHAR_BUDGET
        for msg in useful[-10:]:
            entry = f"{msg.get('sender', '?')}: {msg.get('content', '')[:200]}"
            if len("\n".join(lines + [entry])) > budget:
                break
            lines.append(entry)

        return "\n".join(lines)

    # ------------------------------------------------------------------
    # Layer 4: File windows (large files)
    # ------------------------------------------------------------------

    def _build_file_window(self, target: Path) -> Optional[FileWindow]:
        """Build focused reading window for a file."""
        if not target.exists() or not target.suffix == ".py":
            return None

        try:
            text = target.read_text(encoding="utf-8", errors="replace")
        except Exception:
            return None

        all_lines = text.splitlines()
        if len(all_lines) <= LARGE_FILE_THRESHOLD:
            return None  # small file, no windowing needed

        # Header: first N lines covering imports and module docstring
        header_end = 0
        in_docstring = False
        for i, line in enumerate(all_lines[:60]):
            stripped = line.strip()
            if stripped.startswith(('"""', "'''")):
                in_docstring = not in_docstring
            elif (
                not in_docstring and stripped
                and not stripped.startswith(("#", "import ", "from "))
            ):
                header_end = i
                break
        else:
            header_end = min(60, len(all_lines))

        header = "\n".join(all_lines[:header_end])

        rel = (
            str(target.relative_to(self.root))
            if target.is_relative_to(self.root)
            else str(target)
        )
        return FileWindow(
            path=rel,
            header=header,
            target="",  # target section needs a search term, populated by caller
            total_lines=len(all_lines),
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _resolve(self, filepath: str) -> Optional[Path]:
        """Resolve a manifest path to an absolute path within project root."""
        try:
            p = Path(filepath)
            if not p.is_absolute():
                p = self.root / p
            p = p.resolve()
            # Safety: must be within project root
            if not str(p).startswith(str(self.root)):
                logger.warning("[enrichment] Path escapes project root: %s", filepath)
                return None
            return p
        except Exception:
            return None
