"""Sub-agent: analyze codebase references to refine metadata understanding."""

from __future__ import annotations

import re
from typing import Any

from amx.agents._prompt_helpers import alternatives_mode_directive
from amx.agents.base import (
    AgentContext,
    BaseAgent,
    Confidence,
    MetadataSuggestion,
    apply_logprob_confidence,
)
from amx.codebase.analyzer import CodebaseReport, CodeReference
from amx.config import DEFAULT_ALTERNATIVES_MODE, PromptDetail
from amx.llm.prompts import (
    ALTERNATIVES_LENGTH_RULE_REMINDER,
    length_rule,
    per_col_token_budget,
)
from amx.llm.provider import LLMProvider
from amx.llm.style.guard import scrub_placeholders
from amx.llm.style.injector import render_style_section
from amx.llm.style.loader import load_active_style_profile
from amx.llm.style.profile import StyleProfile
from amx.utils.console import step_spinner
from amx.utils.logging import get_logger
from amx.utils.token_tracker import estimate_tokens, tracker

log = get_logger("agents.code")


def _scrub_suggestions(
    suggestions: list[MetadataSuggestion],
    active_profile: StyleProfile | None,
) -> list[MetadataSuggestion]:
    """Scrub placeholder literals from suggestion text when a style profile is active."""
    if active_profile is None:
        return suggestions
    for s in suggestions:
        s.suggestions = [scrub_placeholders(text) for text in s.suggestions]
    return suggestions


_BASE_SYSTEM_PROMPT = """\
You are a data-catalog expert analyzing how database tables and columns are used
in application code to understand their meaning.

You are given:
- A table in a schema with its list of columns.
- Code snippets where these assets are referenced.

Based on how the code uses these assets, infer a description for EACH column.

Output rules:
- Write every description and reasoning string in **clear, business-friendly American English**.
- Use complete sentences ending with a period.
- Length rule (CRITICAL — honour the user's verbosity preset): {description_length_rule}
- Keep the response labels (`COLUMN`, `DESCRIPTION_1`, `CONFIDENCE`, `REASONING`) verbatim.

Write descriptions assertively and directly (e.g. "Telephone extension number.").
Do NOT start descriptions with "This column likely represents" or "This column is".
Use code behavior as primary evidence: joins, filters, comparisons, assignments, formatting, validations, DTO names, API payload names, and query aliases.
Do not infer meaning from a mere name mention if the snippet does not show behavior.
Never invent a business process, document type, or regulatory meaning unless the code clearly supports it.
If evidence is weak or only the column name appears, give a broader neutral description and lower confidence.
Confidence rules:
- HIGH: behavior is explicit in multiple snippets or naming plus usage strongly agree.
- MEDIUM: the code suggests a likely role but the exact business meaning is not fully proven.
- LOW: only superficial or sparse references exist.
Reasoning must mention the concrete code patterns that justify the description.

{alternatives_mode_directive}
{alternatives_length_reminder}
Respond in this format for each column (one block per column):

COLUMN: <column_name>
DESCRIPTION_1: <best description based on code usage>
{desc_lines}
CONFIDENCE: <HIGH|MEDIUM|LOW>
REASONING: <what code patterns support this>

Example style:
COLUMN: tenant_key
DESCRIPTION_1: Tenant identifier used to scope records.
CONFIDENCE: HIGH
REASONING: The code joins, filters, and groups records by this field across related queries and request payloads.
"""


def _build_system_prompt(
    n_alternatives: int,
    description_verbosity: str = "brief",
    style_profile: StyleProfile | None = None,
    alternatives_mode: str = DEFAULT_ALTERNATIVES_MODE,
) -> str:
    n = max(1, min(5, n_alternatives))
    if n > 1:
        desc_lines = "\n".join(
            f"DESCRIPTION_{i}: <alternative description — apply the SAME length rule as DESCRIPTION_1>"
            for i in range(2, n + 1)
        )
        alternatives_length_reminder = ALTERNATIVES_LENGTH_RULE_REMINDER
    else:
        desc_lines = ""
        alternatives_length_reminder = ""
    return (
        _BASE_SYSTEM_PROMPT.format(
            desc_lines=desc_lines,
            description_length_rule=length_rule(description_verbosity),
            alternatives_mode_directive=alternatives_mode_directive(alternatives_mode, n),
            alternatives_length_reminder=alternatives_length_reminder,
        ).strip()
        + "\n"
        + render_style_section(style_profile)
    )


class CodeAgent(BaseAgent):
    name = "code_agent"

    def __init__(self, llm: LLMProvider, report: CodebaseReport | None = None):
        self.llm = llm
        self.report = report
        self._style_profile = load_active_style_profile()
        # PR γ: snapshot of the retrieval hits + regex refs that
        # informed the most recent ``_build_messages`` call per
        # ``(schema, table)``. The agent's ``run`` / ``parse_batch_result``
        # consult this to build :class:`Citation` records on every
        # produced :class:`MetadataSuggestion` -- matching the
        # ``last_prompt_hits`` contract already shipping in
        # :class:`amx.agents.rag_agent.RAGAgent`.
        self.last_prompt_hits: dict[tuple[str, str], list[dict]] = {}
        self.last_prompt_refs: dict[tuple[str, str], list[Any]] = {}
        # PR δ (I1): diagnostics buffer mirrored from RAGAgent. The
        # orchestrator's ``table_processor`` drains it after each
        # table to surface user-relevant signals that aren't fatal
        # (empty context, query timeout, etc.) so the user sees
        # why this table produced zero code suggestions instead of a
        # silent skip.
        self._diagnostics: list[str] = []

    def consume_diagnostics(self) -> list[str]:
        """Return + clear the diagnostics buffer (PR δ — mirrors RAGAgent)."""
        diagnostics = list(self._diagnostics)
        self._diagnostics.clear()
        return diagnostics

    def _record_diagnostic(self, message: str) -> None:
        if message:
            self._diagnostics.append(message)

    def _code_query_timeout(self) -> float | None:
        """Resolved per-query timeout for code retrieval (seconds), or ``None``.

        Re-uses :attr:`LLMConfig.rag_query_timeout_sec` — one timeout
        knob covers both code and docs since both go through Chroma.
        """
        raw = getattr(self.llm.cfg, "rag_query_timeout_sec", 5.0)
        try:
            value = float(raw)
        except (TypeError, ValueError):
            return None
        return value if value > 0 else None

    @property
    def _n_alternatives(self) -> int:
        return max(1, min(5, getattr(self.llm.cfg, "n_alternatives", 3)))

    @property
    def _description_verbosity(self) -> str:
        return getattr(self.llm.cfg, "description_verbosity", "brief")

    def _per_col_token_budget(self) -> int:
        return per_col_token_budget(self._description_verbosity)

    @property
    def _prompt_detail(self) -> PromptDetail:
        # Defensive fallback for tests / call sites that pass a minimal
        # LLM stub without ``prompt_detail_cfg``: return the default
        # preset so ``_build_messages`` keeps the pre-PR-γ behaviour
        # (``code_col_hits=0``, single table-level semantic query).
        cfg = getattr(self.llm, "cfg", None)
        pd = getattr(cfg, "prompt_detail_cfg", None)
        if pd is not None:
            return pd
        return PromptDetail()

    def _build_messages(self, ctx: AgentContext) -> list[dict[str, str]] | None:
        """Build the Code Agent prompt messages. Returns ``None`` when no code context exists."""
        from amx.codebase.code_rag import code_collection_count, query_code_snippets
        from amx.docs.rag import RAGQueryTimeout

        # PR δ (C8): re-run snapshot path — when the orchestrator
        # hydrated the AgentContext from a stored payload that already
        # carries the exact code chunks the original run consumed, skip
        # the live Chroma queries entirely. Makes re-runs deterministic
        # across re-indexes and avoids re-paying retrieval latency on
        # every re-run.
        snapshot_hits = list(getattr(ctx, "code_hits", None) or [])

        if not self.report and not snapshot_hits:
            return None

        has_refs = bool(self.report and (self.report.references or self.report.external_mentions))
        source_filters = (
            [p for p in self.report.path.split(";") if p]
            if (self.report and self.report.path)
            else None
        )

        if snapshot_hits:
            has_sem = True
        else:
            try:
                has_sem = code_collection_count(source_filters=source_filters) > 0
            except Exception:
                has_sem = False
        if not has_refs and not has_sem:
            return None

        columns = ctx.db_profile.get("columns", [])
        if not columns:
            return None

        pd = self._prompt_detail

        table_refs = (
            self.report.references.get(ctx.table.lower(), [])
            if (self.report and self.report.references)
            else []
        )

        all_code_blocks: list[str] = []
        consumed_refs: list[CodeReference] = []

        if table_refs:
            all_code_blocks.append(
                "## Table-level references\n"
                + "\n---\n".join(f"File: {r.file}:{r.line_no}\n{r.context}" for r in table_refs[:8])
            )
            consumed_refs.extend(table_refs[:8])

        if self.report and self.report.references:
            for col in columns:
                col_name = col["name"].lower()
                refs = self.report.references.get(col_name, [])
                if refs:
                    all_code_blocks.append(
                        f"## Column: {col['name']}\n"
                        + "\n---\n".join(
                            f"File: {r.file}:{r.line_no}\n{r.context}" for r in refs[:5]
                        )
                    )
                    consumed_refs.extend(refs[:5])

        # PR γ: track every semantic hit fed into the prompt so we can
        # turn them into :class:`Citation` records after the LLM answers.
        sem_hits_all: list[dict] = []
        any_timed_out = False

        if snapshot_hits:
            sem_hits_all.extend(snapshot_hits)
        elif has_sem:
            timeout = self._code_query_timeout()
            # PR γ: drop the hardcoded "SQL Spark dataframe usage" bias
            # — it skewed retrieval toward Spark / PySpark snippets even
            # for repos that don't use Spark. A neutral
            # ``"<schema> <table>"`` query is the cheapest correct default.
            table_query = f"{ctx.schema} {ctx.table}".strip()
            try:
                table_hits = query_code_snippets(
                    table_query,
                    n_results=5,
                    source_filters=source_filters,
                    timeout=timeout,
                )
                sem_hits_all.extend(table_hits)
            except RAGQueryTimeout:
                any_timed_out = True

            # PR γ: optional per-column fan-out, parallel to the
            # ``rag_col_hits`` knob on docs-RAG. Default ``0`` keeps the
            # previous behaviour byte-for-byte; ``detailed`` / ``full``
            # raise it to ``1`` so column call-sites surface.
            col_hit_limit = int(getattr(pd, "code_col_hits", 0) or 0)
            if col_hit_limit > 0 and not any_timed_out:
                seen_keys: set[str] = {h["text"][:120] for h in sem_hits_all}
                for col in columns:
                    try:
                        col_hits = query_code_snippets(
                            f"{ctx.table}.{col['name']} column",
                            n_results=col_hit_limit,
                            source_filters=source_filters,
                            timeout=timeout,
                        )
                    except RAGQueryTimeout:
                        any_timed_out = True
                        break
                    for h in col_hits:
                        key = h["text"][:120]
                        if key in seen_keys:
                            continue
                        seen_keys.add(key)
                        sem_hits_all.append(h)

        if any_timed_out:
            timeout_val = self._code_query_timeout() or 0.0
            self._record_diagnostic(
                f"Code: retrieval timed out after {timeout_val:.1f}s for "
                f"{ctx.schema}.{ctx.table}; falling back to regex refs only"
            )

        if sem_hits_all:
            all_code_blocks.append(
                "## Semantic code retrieval (nearest chunks)\n"
                + "\n---\n".join(h["text"][:900] for h in sem_hits_all)
            )

        ext_flat: list[CodeReference] = []
        if self.report:
            for lst in (self.report.external_mentions or {}).values():
                ext_flat.extend(lst[:2])
        if ext_flat:
            all_code_blocks.append(
                "## Other identifiers (not in connected DB catalog)\n"
                + "\n---\n".join(f"File: {r.file}:{r.line_no}\n{r.context}" for r in ext_flat[:5])
            )
            consumed_refs.extend(ext_flat[:5])

        if not all_code_blocks:
            return None

        # Stash hits + regex refs so the agent's caller (``run`` /
        # ``parse_batch_result``) can attach citations to every emitted
        # suggestion. Keyed by ``(schema, table)`` because one CodeAgent
        # instance handles many tables in a single run.
        self.last_prompt_hits[(ctx.schema, ctx.table)] = list(sem_hits_all)
        self.last_prompt_refs[(ctx.schema, ctx.table)] = list(consumed_refs)

        from amx.agents.base import _user_instructions_block

        col_lines = "\n".join(f"  - {c['name']} (type={c['dtype']})" for c in columns)
        user_msg = (
            f"Schema: {ctx.schema}\n"
            f"Table: {ctx.table}\n\n"
            f"Columns:\n{col_lines}\n\n"
            f"Code references:\n\n" + "\n\n".join(all_code_blocks) + _user_instructions_block(ctx)
        )
        system = _build_system_prompt(
            self._n_alternatives,
            self._description_verbosity,
            style_profile=self._style_profile,
            alternatives_mode=getattr(self.llm.cfg, "alternatives_mode", DEFAULT_ALTERNATIVES_MODE),
        )
        return [
            {"role": "system", "content": system},
            {"role": "user", "content": user_msg},
        ]

    def _scaled_max_tokens(self, n_columns: int) -> int:
        return max(self.llm.cfg.max_tokens, n_columns * self._per_col_token_budget())

    def collect_messages(self, ctx: AgentContext) -> list:
        """Return a ``BatchRequest`` for this table (or empty list when no code context)."""
        from amx.llm.batch import BatchRequest

        msgs = self._build_messages(ctx)
        if msgs is None:
            self._record_diagnostic(
                f"Code: no codebase context (regex refs or semantic hits) found "
                f"for {ctx.schema}.{ctx.table}"
            )
            return []
        n_columns = len(ctx.db_profile.get("columns", []) or [])
        return [
            BatchRequest(
                custom_id=f"code:{ctx.schema}:{ctx.table}",
                messages=msgs,
                max_tokens=self._scaled_max_tokens(n_columns),
                temperature=self.llm.cfg.temperature,
                metadata={"schema": ctx.schema, "table": ctx.table},
            )
        ]

    def _attach_code_citations(
        self,
        suggestions: list[MetadataSuggestion],
        ctx: AgentContext,
    ) -> list[MetadataSuggestion]:
        """Build :class:`Citation` records from the stashed hits + refs.

        Mirrors :func:`amx.agents.rag_agent._attach_citations`. Both the
        regex-references path and the semantic-retrieval path produce
        citations; they're deduped together by ``(source, chunk_idx,
        line_range)`` so a snippet surfaced through both channels
        renders once.
        """
        from amx.agents._citations import (
            attach_citations,
            hits_to_citations,
            regex_refs_to_citations,
        )

        if not suggestions:
            return suggestions
        key = (ctx.schema, ctx.table)
        hits = self.last_prompt_hits.get(key) or []
        refs = self.last_prompt_refs.get(key) or []

        citations = hits_to_citations(hits)
        # Re-pack ``[CodeReference]`` into the ``dict[asset -> refs]``
        # shape :func:`regex_refs_to_citations` expects. The dedup
        # inside that helper drops the duplicates per ``(file, line_no)``.
        if refs:
            citations = list(citations) + list(regex_refs_to_citations({"_": refs}))
        return attach_citations(suggestions, citations)

    def parse_batch_result(self, content: str, ctx: AgentContext) -> list[MetadataSuggestion]:
        """Parse a raw LLM text response; used after Batch API completes."""
        suggestions = self._parse_response(content, ctx)
        suggestions = _scrub_suggestions(suggestions, self._style_profile)
        return self._attach_code_citations(suggestions, ctx)

    def run(self, ctx: AgentContext) -> list[MetadataSuggestion]:
        self._diagnostics.clear()
        messages = self._build_messages(ctx)
        if messages is None:
            log.info("No code context for %s.%s, skipping", ctx.schema, ctx.table)
            self._record_diagnostic(
                f"Code: no codebase context (regex refs or semantic hits) found "
                f"for {ctx.schema}.{ctx.table}"
            )
            return []

        columns = ctx.db_profile.get("columns", [])
        est = estimate_tokens(messages)
        mt = self._scaled_max_tokens(len(columns))
        with step_spinner(f"Code Agent: {len(columns)} columns", token_estimate=est):
            result = self.llm.chat(messages, max_tokens=mt)
        tracker.record_for("code_agent", est, self.llm, result.usage)

        suggestions = self._parse_response(result.content, ctx)
        suggestions = _scrub_suggestions(suggestions, self._style_profile)
        suggestions = apply_logprob_confidence(
            suggestions,
            result.logprobs,
            high_threshold=self.llm.cfg.logprob_high,
            medium_threshold=self.llm.cfg.logprob_medium,
            response_text=result.content,
        )
        return self._attach_code_citations(suggestions, ctx)

    def _parse_response(
        self, text: str, ctx: AgentContext, default_col: str = ""
    ) -> list[MetadataSuggestion]:
        text = re.sub(r"^```[a-zA-Z0-9_-]*\s*", "", (text or "").strip())
        text = re.sub(r"\s*```$", "", text).strip()
        suggestions: list[MetadataSuggestion] = []
        current_col = default_col
        descs: list[str] = []
        conf = Confidence.MEDIUM
        reasoning = ""

        for line in text.splitlines():
            line = line.strip()
            if line.startswith("COLUMN:"):
                if current_col and descs:
                    suggestions.append(
                        MetadataSuggestion(
                            schema=ctx.schema,
                            table=ctx.table,
                            column=current_col,
                            suggestions=descs,
                            confidence=conf,
                            reasoning=reasoning,
                            source="codebase",
                        )
                    )
                current_col = line.split(":", 1)[1].strip()
                descs = []
                conf = Confidence.MEDIUM
                reasoning = ""
            elif line.startswith("DESCRIPTION_"):
                descs.append(line.split(":", 1)[1].strip())
            elif line.startswith("CONFIDENCE:"):
                raw = line.split(":", 1)[1].strip().upper()
                conf = Confidence[raw] if raw in Confidence.__members__ else Confidence.MEDIUM
            elif line.startswith("REASONING:"):
                reasoning = line.split(":", 1)[1].strip()

        if current_col and descs:
            suggestions.append(
                MetadataSuggestion(
                    schema=ctx.schema,
                    table=ctx.table,
                    column=current_col,
                    suggestions=descs,
                    confidence=conf,
                    reasoning=reasoning,
                    source="codebase",
                )
            )

        return suggestions
