"""Unified context assembler for cross-system queries.

Takes a query and builds an optimal LLM context window by searching across
code, docs, and telemetry subsystems in parallel.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class CodeSnippet:
    """A ranked code search result."""
    entity_id: int
    name: str
    qualified_name: str
    file_path: str
    start_line: int
    end_line: int
    code: str
    language: str
    score: float
    version: str = ""


@dataclass(slots=True)
class DocChunk:
    """A ranked document chunk result."""
    chunk_id: int
    document_id: int
    title: str
    content: str
    score: float
    version: int = 1
    is_latest: bool = True


@dataclass(slots=True)
class TelemetryItem:
    """A relevant telemetry result (log, trace, or metric)."""
    item_type: str  # "log", "trace", "metric"
    service: str
    summary: str
    timestamp: str
    severity: str = ""
    score: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class AssembledContext:
    """The assembled cross-system context for an LLM query."""
    query: str
    code_snippets: list[CodeSnippet] = field(default_factory=list)
    doc_chunks: list[DocChunk] = field(default_factory=list)
    telemetry: list[TelemetryItem] = field(default_factory=list)
    graph_context: dict[str, Any] = field(default_factory=dict)
    token_count: int = 0
    sources: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        from dataclasses import asdict
        return asdict(self)

    def to_prompt_text(self) -> str:
        """Format as a text block suitable for LLM context injection."""
        parts = []
        if self.code_snippets:
            parts.append("## Relevant Code\n")
            for s in self.code_snippets:
                parts.append(f"### {s.qualified_name or s.name} ({s.file_path}:{s.start_line}-{s.end_line})")
                parts.append(f"```{s.language}\n{s.code}\n```\n")
        if self.doc_chunks:
            parts.append("## Relevant Documentation\n")
            for d in self.doc_chunks:
                parts.append(f"### {d.title} (score: {d.score:.2f})")
                parts.append(f"{d.content}\n")
        if self.telemetry:
            parts.append("## Relevant Telemetry\n")
            for t in self.telemetry:
                parts.append(f"- [{t.item_type}] {t.service} @ {t.timestamp}: {t.summary}")
        if self.graph_context:
            parts.append("## Entity Relationships\n")
            nodes = self.graph_context.get("nodes", [])
            edges = self.graph_context.get("edges", [])
            if nodes:
                parts.append(f"Entities: {', '.join(n.get('name', '') for n in nodes[:20])}")
            if edges:
                parts.append(f"Relationships: {len(edges)} connections")
        return "\n".join(parts)


def _estimate_tokens(text: str) -> int:
    """Rough token estimate (1 token ~ 4 chars)."""
    return max(1, len(text) // 4)


def route_query(query: str) -> list[str]:
    """Determine which subsystems are relevant to a query.

    Returns list of subsystem names: "code", "docs", "otel".
    """
    q = query.lower()
    subsystems = []

    code_signals = ["function", "class", "method", "code", "import", "module",
                    "file", "entity", "dependency", "call", "implement", "def ",
                    "source", "refactor", "bug", "error in"]
    doc_signals = ["document", "docs", "explain", "what is", "how does",
                   "architecture", "design", "policy", "guide", "spec",
                   "requirement", "billing", "payment"]
    otel_signals = ["trace", "log", "metric", "latency", "error rate", "slow",
                    "timeout", "service", "health", "alert", "incident",
                    "cpu", "memory", "disk", "deployment", "restart", "crash"]

    if any(s in q for s in code_signals):
        subsystems.append("code")
    if any(s in q for s in doc_signals):
        subsystems.append("docs")
    if any(s in q for s in otel_signals):
        subsystems.append("otel")

    # If no clear signal, search everything
    if not subsystems:
        subsystems = ["code", "docs", "otel"]

    return subsystems


class ContextAssembler:
    """Assembles cross-system context for agentic AI queries.

    Searches code (via ChromaDB), docs (via ChromaDB), and telemetry
    (via SQLite), then ranks, deduplicates, and truncates to fit
    a token budget.
    """

    def __init__(
        self,
        code_repo=None,
        docs_repo=None,
        otel_repo=None,
        code_vectorstore=None,
        docs_vectorstore=None,
        embedder=None,
        graph_store=None,
    ) -> None:
        self.code_repo = code_repo
        self.docs_repo = docs_repo
        self.otel_repo = otel_repo
        self.code_vectorstore = code_vectorstore
        self.docs_vectorstore = docs_vectorstore
        self.embedder = embedder
        self.graph_store = graph_store

    def assemble(
        self,
        query: str,
        context_budget: int = 8000,
        subsystems: list[str] | None = None,
        code_version: str | None = None,
    ) -> AssembledContext:
        """Build an optimal LLM context window for the query.

        1. Route query to relevant subsystems
        2. Search each subsystem
        3. Graph expansion for found entities
        4. Rank all results by relevance
        5. Truncate to fit context_budget (token count)
        6. Return structured context with source attribution
        """
        ctx = AssembledContext(query=query)
        targets = subsystems or route_query(query)
        ctx.metadata["subsystems_queried"] = targets

        # Embed the query once
        query_embedding = None
        if self.embedder:
            try:
                query_embedding = self.embedder.embed_query(query)
            except Exception as exc:
                logger.warning("Failed to embed query: %s", exc)

        # Search each subsystem
        if "code" in targets:
            ctx.code_snippets = self._search_code(query, query_embedding, code_version)
            ctx.sources.extend(f"code:{s.file_path}" for s in ctx.code_snippets)

        if "docs" in targets:
            ctx.doc_chunks = self._search_docs(query, query_embedding)
            ctx.sources.extend(f"docs:{d.title}" for d in ctx.doc_chunks)

        if "otel" in targets:
            ctx.telemetry = self._search_telemetry(query)
            ctx.sources.extend(f"otel:{t.service}" for t in ctx.telemetry)

        # Graph expansion
        if self.graph_store:
            entity_names = [s.qualified_name or s.name for s in ctx.code_snippets[:5]]
            entity_names.extend(d.title for d in ctx.doc_chunks[:3])
            for name in entity_names:
                if self.graph_store.has_node(name):
                    neighborhood = self.graph_store.neighborhood(name, depth=1, max_nodes=20)
                    if neighborhood.get("nodes"):
                        ctx.graph_context = neighborhood
                        break

        # Truncate to budget
        ctx = self._truncate_to_budget(ctx, context_budget)
        ctx.token_count = _estimate_tokens(ctx.to_prompt_text())

        return ctx

    def _search_code(
        self, query: str, query_embedding: list[float] | None, version: str | None = None,
    ) -> list[CodeSnippet]:
        """Search code via ChromaDB vector search, falling back to text search."""
        results = []

        # ChromaDB vector search
        if self.code_vectorstore and query_embedding:
            hits = self.code_vectorstore.search_code(query_embedding, top_k=10, version=version)
            for hit in hits:
                meta = hit.get("metadata", {})
                results.append(CodeSnippet(
                    entity_id=hit["entity_id"],
                    name=meta.get("name", ""),
                    qualified_name=meta.get("qualified_name", ""),
                    file_path=meta.get("file_path", ""),
                    start_line=int(meta.get("start_line", 0)),
                    end_line=int(meta.get("end_line", 0)),
                    code=hit.get("code", ""),
                    language=meta.get("language", ""),
                    score=hit["score"],
                    version=version or "",
                ))

        # Fallback: text search via repository
        if not results and self.code_repo:
            text_results = self.code_repo.search_entities(query, limit=10)
            for r in text_results:
                results.append(CodeSnippet(
                    entity_id=r.get("id", 0),
                    name=r.get("name", ""),
                    qualified_name=r.get("qualified_name", ""),
                    file_path=str(r.get("file_id", "")),
                    start_line=r.get("start_line", 0),
                    end_line=r.get("end_line", 0),
                    code="",
                    language=r.get("language", ""),
                    score=0.5,
                ))

        return results

    def _search_docs(
        self, query: str, query_embedding: list[float] | None,
    ) -> list[DocChunk]:
        """Search docs via ChromaDB, falling back to repository search."""
        results = []

        if self.docs_vectorstore and query_embedding:
            hits = self.docs_vectorstore.search_chunks(query_embedding, top_k=10)
            for hit in hits:
                meta = hit.get("metadata", {})
                results.append(DocChunk(
                    chunk_id=hit["chunk_id"],
                    document_id=int(meta.get("document_id", 0)),
                    title=meta.get("title", ""),
                    content=hit.get("content", ""),
                    score=hit["score"],
                ))

        # Fallback: repository keyword search
        if not results and self.docs_repo:
            try:
                search_results = self.docs_repo.search(query, limit=10)
                for r in search_results:
                    results.append(DocChunk(
                        chunk_id=r.get("chunk_id", 0),
                        document_id=r.get("document_id", 0),
                        title=r.get("title", ""),
                        content=r.get("content", ""),
                        score=r.get("score", 0.5),
                    ))
            except Exception:
                pass

        return results

    def _search_telemetry(self, query: str) -> list[TelemetryItem]:
        """Search telemetry via OtelRepository."""
        results = []
        if not self.otel_repo:
            return results

        try:
            # Search logs
            log_results = self.otel_repo.search_logs(query, limit=5)
            for log in log_results:
                results.append(TelemetryItem(
                    item_type="log",
                    service=log.get("service_name", "unknown"),
                    summary=log.get("body", "")[:500],
                    timestamp=str(log.get("timestamp", "")),
                    severity=log.get("severity", ""),
                    score=0.5,
                ))
        except Exception:
            pass

        try:
            # Search traces
            trace_results = self.otel_repo.search_traces(query, limit=5)
            for trace in trace_results:
                results.append(TelemetryItem(
                    item_type="trace",
                    service=trace.get("root_service_name", "unknown"),
                    summary=f"{trace.get('root_span_name', '')} ({trace.get('status_code', '')})",
                    timestamp=str(trace.get("start_time", "")),
                    score=0.5,
                    metadata={"trace_id": trace.get("trace_id", ""), "duration_ms": trace.get("duration_ms", 0)},
                ))
        except Exception:
            pass

        return results

    def _truncate_to_budget(self, ctx: AssembledContext, budget: int) -> AssembledContext:
        """Truncate context to fit within token budget, keeping highest-scored items."""
        total = _estimate_tokens(ctx.to_prompt_text())
        if total <= budget:
            return ctx

        # Progressively drop lowest-scored items
        while _estimate_tokens(ctx.to_prompt_text()) > budget:
            # Drop lowest-scored telemetry first, then docs, then code
            if ctx.telemetry:
                ctx.telemetry.pop()
            elif ctx.doc_chunks:
                ctx.doc_chunks.pop()
            elif ctx.code_snippets:
                ctx.code_snippets.pop()
            else:
                break

        return ctx
