src/agents/report/generator.py

"""ReportGenerator — turns a session's AnalysisRecords into an AnalysisReport (KM-644).

A button-triggered service shaped like the Assembler: deterministic assembly of the
records (findings/caveats/open_questions/data_sources/method_steps, copied verbatim —
INV-4) wrapped around exactly ONE LLM call that authors only the executive summary.
If that call fails the report is still returned with a deterministic fallback
summary (decision D1) — the deterministic body is the real value.

Versioning + persistence live in `ReportStore`; this service does generation only
(returns an `AnalysisReport` with `version=0`; the store assigns the real version).
Chain construction mirrors `agents/slow_path/assembler.py`.
"""

from __future__ import annotations

import re
from datetime import UTC, datetime
from pathlib import Path

from langchain_core.messages import SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from langchain_openai import AzureChatOpenAI

from src.middlewares.logging import get_logger

from ..language import detect_reply_language
from ..slow_path.schemas import AnalysisRecord, TaskSummary
from .errors import ReportError
from .readiness import has_successful_analysis
from .schemas import (
    AnalysisReport,
    AttributedNote,
    BQAnswerDraft,
    BusinessQuestionAnswer,
    DataSourceRef,
    EvidenceTable,
    ProblemStatement,
    ReportFinding,
    ReportSummaryNarrative,
)

logger = get_logger("report_generator")

_FALLBACK_SUMMARY = "Automated summary unavailable — see the findings below."

# Caps keeping the deterministic sections readable on multi-record analyses.
_MAX_CAVEATS = 12
_MAX_OPEN_QUESTIONS = 10
_EVIDENCE_MAX_ROWS = 10
_EVIDENCE_MAX_TABLES = 3  # per record
# Wider tables are raw analysis *inputs* (e.g. a 19-column correlation pull), not
# presentable evidence — grouped/top-N/merge results are always narrow.
_EVIDENCE_MAX_COLS = 8

# CRISP-DM phases in narrative order, with human labels for the method appendix.
_STAGE_LABELS: list[tuple[str, str]] = [
    ("data_understanding", "Data understanding"),
    ("data_preparation", "Data preparation"),
    ("modeling", "Modeling"),
    ("evaluation", "Evaluation"),
]

# Human labels for BusinessQuestionAnswer.status in the rendered markdown.
_BQ_STATUS_LABELS: dict[str, str] = {
    "answered": "Answered",
    "partial": "Partially answered",
    "unanswered": "Unanswered",
}

# Friendly labels for the catalog's internal source_type enum, shown in Data Sources.
_SOURCE_TYPE_LABELS: dict[str, str] = {
    "schema": "Database",
    "tabular": "Tabular file",
    "unstructured": "Documents",
}

_PROMPT_PATH = (
    Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "report_summary.md"
)


def _load_prompt_text() -> str:
    return _PROMPT_PATH.read_text(encoding="utf-8")


def _build_default_chain() -> Runnable:
    from src.config.settings import settings

    llm = AzureChatOpenAI(
        azure_deployment=settings.azureai_deployment_name_4o,
        openai_api_version=settings.azureai_api_version_4o,
        azure_endpoint=settings.azureai_endpoint_url_4o,
        api_key=settings.azureai_api_key_4o,
        temperature=0,
    )
    prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(content=_load_prompt_text()),
            ("human", "{human_content}"),
        ]
    )
    return prompt | llm.with_structured_output(ReportSummaryNarrative)


_default_chain: Runnable | None = None


def _get_default_chain() -> Runnable:
    global _default_chain
    if _default_chain is None:
        _default_chain = _build_default_chain()
    return _default_chain


# --------------------------------------------------------------------------- #
# Deterministic assembly (pure; no LLM, no I/O) — easy to unit-test.
# --------------------------------------------------------------------------- #


def _collect_findings(records: list[AnalysisRecord]) -> list[ReportFinding]:
    # Each finding traces to its record. Deduped *within* a record (an Assembler run
    # occasionally repeats a line); kept across records so the grouped render can show
    # each analysis's own findings under its own question.
    out: list[ReportFinding] = []
    for rec in records:
        seen: set[str] = set()
        for text in rec.findings:
            if text in seen:
                continue
            seen.add(text)
            out.append(ReportFinding(text=text, record_ids=[rec.record_id]))
    return out


def _note_key(text: str) -> str:
    # Dedupe key: collapse whitespace, drop trailing punctuation, casefold — so the
    # Assembler's near-identical rephrasings ("Data is capped at 500 rows." vs
    # "data is capped at 500 rows") merge into one note.
    return " ".join(text.split()).rstrip(".!").casefold()


def _collect_notes(records: list[AnalysisRecord], field: str, cap: int) -> list[AttributedNote]:
    # Caveats / open_questions are deduped by normalized text; a merged note keeps
    # the first phrasing seen and cites every record it came from (plural
    # record_ids). Capped so a many-record analysis stays readable.
    merged: dict[str, AttributedNote] = {}
    for rec in records:
        for text in getattr(rec, field):
            key = _note_key(text)
            if not key:
                continue
            note = merged.setdefault(key, AttributedNote(text=text))
            if rec.record_id not in note.record_ids:
                note.record_ids.append(rec.record_id)
    return list(merged.values())[:cap]


def _fmt_cell(value) -> str:
    if value is None:
        return "—"
    if isinstance(value, float):
        return f"{value:g}"  # 1234.5 not 1234.5000000001; no trailing zeros
    return str(value)


def _collect_evidence(records: list[AnalysisRecord]) -> dict[str, list[EvidenceTable]]:
    """Copy small result tables out of each record's `results_snapshot` (INV-4).

    Table-kind tool outputs only — the copy-paste-able slices (top-N rankings,
    grouped aggregates, merges). `check_*` outputs are skipped (catalog metadata,
    not evidence). Rows and tables-per-record are capped so a wide retrieval
    can't balloon the report.
    """
    out: dict[str, list[EvidenceTable]] = {}
    for rec in records:
        tables: list[EvidenceTable] = []
        for result in rec.results_snapshot.values():
            for output in result.outputs:
                if len(tables) >= _EVIDENCE_MAX_TABLES:
                    break
                if output.tool in ("check_data", "check_knowledge"):
                    continue
                if output.kind != "table" or not output.columns or not output.rows:
                    continue
                if len(output.columns) > _EVIDENCE_MAX_COLS:
                    continue
                tables.append(
                    EvidenceTable(
                        title=result.objective,
                        columns=[str(c) for c in output.columns],
                        rows=[
                            [_fmt_cell(v) for v in row]
                            for row in output.rows[:_EVIDENCE_MAX_ROWS]
                        ],
                        truncated=len(output.rows) > _EVIDENCE_MAX_ROWS,
                    )
                )
        if tables:
            out[rec.record_id] = tables
    return out


def _unresolved_note(rec: AnalysisRecord) -> AttributedNote:
    # Goal + the record's own first caveat as the "why" — both Assembler-authored,
    # nothing new is synthesized here.
    text = rec.goal_restated or "Analysis run"
    reason = next(iter(rec.caveats), None)
    if reason:
        text += f" — {reason}"
    return AttributedNote(text=text, record_ids=[rec.record_id])


def _excluded_note(rec: AnalysisRecord) -> AttributedNote:
    return AttributedNote(
        text=rec.goal_restated or rec.record_id, record_ids=[rec.record_id]
    )


def _collect_method_steps(records: list[AnalysisRecord]) -> list[TaskSummary]:
    steps: list[TaskSummary] = []
    for rec in records:
        steps.extend(rec.tasks_run)
    return steps


def _build_data_sources(
    records: list[AnalysisRecord], catalog
) -> list[DataSourceRef]:
    """Freeze real catalog metadata for the sources this analysis used.

    `catalog` is the analysis-scope catalog — already restricted to this analysis's
    bound sources — so every source in it is a candidate. Matches candidates against
    the records' (narrative) `data_used` by name/id; falls back to all sources, then
    to bare `data_used` strings if no catalog is available — so the section is always
    populated, best-effort.
    """
    if catalog is None or not catalog.sources:
        seen: list[str] = []
        for rec in records:
            for du in rec.data_used:
                if du not in seen:
                    seen.append(du)
        return [DataSourceRef(source_id=d, name=d, source_type="", detail={}) for d in seen]

    candidates = catalog.sources

    def _ref(s) -> DataSourceRef:
        return DataSourceRef(
            source_id=s.source_id,
            name=s.name,
            source_type=s.source_type,
            detail={
                "tables": [t.name for t in s.tables],
                "row_count": sum((t.row_count or 0) for t in s.tables) or None,
                "columns": [c.name for t in s.tables for c in t.columns],
            },
        )

    used = " ".join(du for rec in records for du in rec.data_used).lower()
    matched = [
        _ref(s)
        for s in candidates
        if s.name.lower() in used or s.source_id.lower() in used
    ]
    return matched or [_ref(s) for s in candidates]


def _build_human_content(
    ps: ProblemStatement,
    records: list[AnalysisRecord],
    caveats: list[AttributedNote],
    reply_language: str,
) -> str:
    # Questions and analyses are NUMBERED so the model can reference them by index
    # in `bq_answers` (question_index / analysis_indexes) — it never reproduces ids.
    sections = []
    if ps.objective:
        sections.append("# Objective\n" + ps.objective)
    if ps.business_questions:
        sections.append(
            "# Business questions\n"
            + "\n".join(f"{i}. {q}" for i, q in enumerate(ps.business_questions, 1))
        )
    lines = ["# Analyses (findings already finalized — synthesize, do not add numbers)"]
    for i, rec in enumerate(records, 1):
        lines.append(f"Analysis {i}: {rec.goal_restated}")
        seen: set[str] = set()
        for text in rec.findings:
            if text in seen:
                continue
            seen.add(text)
            lines.append(f"- {text}")
    sections.append("\n".join(lines))
    if caveats:
        sections.append("# Caveats\n" + "\n".join(f"- {c.text}" for c in caveats))
    sections.append("# Reply language\n" + reply_language)
    return "\n\n".join(sections)


def _resolve_bq_answers(
    drafts: list[BQAnswerDraft],
    questions: list[str],
    records: list[AnalysisRecord],
) -> list[BusinessQuestionAnswer]:
    """Map the LLM's index-based drafts onto real question text and record ids.

    Every question gets a row (unanswered when the model skipped it);
    out-of-range indexes are silently dropped.
    """
    if not questions:
        return []
    by_index = {d.question_index: d for d in drafts}
    out: list[BusinessQuestionAnswer] = []
    for i, question in enumerate(questions, 1):
        draft = by_index.get(i)
        if draft is None:
            out.append(BusinessQuestionAnswer(question=question))
            continue
        record_ids = [
            records[j - 1].record_id
            for j in draft.analysis_indexes
            if 1 <= j <= len(records)
        ]
        out.append(
            BusinessQuestionAnswer(
                question=question,
                answer=draft.answer,
                status=draft.status,
                record_ids=record_ids,
            )
        )
    return out


# Inline code spans (one or more backticks). Content inside is already literal in
# Markdown/MDX, so escaping within them would only surface a visible backslash
# (e.g. `product\_id`). We keep code spans verbatim and escape only around them.
_CODE_SPAN_RE = re.compile(r"`+[^`]*`+")


def _mdx_escape(text: str) -> str:
    """Escape Markdown/MDX syntax characters in a dynamic value.

    Applied only to values authored outside the renderer (LLM findings/caveats,
    user objective/questions, catalog source & table names, tool names) so that
    identifiers like ``XL_S_129`` don't italicize and stray ``<``/``{`` don't break
    MDX compilation. Content inside inline code spans is left verbatim (already
    literal). NOT applied to the executive summary (its prompt allows light inline
    emphasis) nor to the structural markdown the renderer emits itself.
    """
    if not text:
        return text

    def _esc(s: str) -> str:
        s = s.replace("\\", "\\\\")
        for ch in ("<", "{", "|", "_", "*"):
            s = s.replace(ch, "\\" + ch)
        return s

    out: list[str] = []
    pos = 0
    for m in _CODE_SPAN_RE.finditer(text):
        out.append(_esc(text[pos : m.start()]))
        out.append(m.group(0))  # keep the code span verbatim
        pos = m.end()
    out.append(_esc(text[pos:]))
    return "".join(out)


def _render_markdown(report: AnalysisReport) -> str:
    # Version is deliberately NOT in the markdown — it is assigned by the store
    # after rendering and lives in the structured `version` field / API metadata.
    meta = f"*Generated {report.generated_at:%Y-%m-%d}"
    author = report.user_name or report.user_id
    if author:
        meta += f" by {author}"
    meta += f" · {len(report.record_ids)} analyses · {len(report.data_sources)} source(s)*"
    # Title + meta form the header block; each subsequent section is divided by a
    # horizontal rule (`---`) so the report reads as a formal, sectioned document.
    parts: list[str] = ["# Analysis Report\n" + meta]

    ps = report.problem_statement
    if ps.objective:
        parts.append("## Objective\n" + _mdx_escape(ps.objective))
    if ps.business_questions:
        parts.append(
            "## Business Questions\n"
            + "\n".join(
                f"{i}. {_mdx_escape(q)}" for i, q in enumerate(ps.business_questions, 1)
            )
        )

    if report.executive_summary:
        parts.append("## Executive Summary\n" + report.executive_summary)

    if report.bq_answers:
        lines = ["## Answers to Business Questions"]
        for i, a in enumerate(report.bq_answers, 1):
            label = _BQ_STATUS_LABELS.get(a.status, a.status)
            entry = f"{i}. **{_mdx_escape(a.question)}** — *{label}*"
            if a.answer:
                # LLM prose (same authorship as the executive summary): not escaped.
                entry += f"\n   {a.answer}"
            lines.append(entry)
        parts.append("\n".join(lines))

    if report.findings:
        # Group findings by their originating analysis (record) so results from
        # different questions read as separate analyses, not one flat, seemingly
        # contradictory list. Subheadings (the restated question) appear only when
        # more than one analysis contributed.
        by_record: dict[str, list[ReportFinding]] = {}
        for f in report.findings:
            by_record.setdefault(f.record_ids[0] if f.record_ids else "", []).append(f)
        ordered = [rid for rid in report.record_goals if rid in by_record]
        ordered += [rid for rid in by_record if rid not in ordered]
        grouped = sum(1 for rid in ordered if by_record.get(rid)) > 1

        blocks = ["## Key Findings"]
        for rid in ordered:
            group = by_record.get(rid)
            if not group:
                continue
            block: list[str] = []
            if grouped:
                block.append(f"### {_mdx_escape(report.record_goals.get(rid) or 'Analysis')}")
            block.extend(f"{i}. {_mdx_escape(f.text)}" for i, f in enumerate(group, 1))
            # Evidence tables (copied result slices) under the findings they back,
            # so the numbers are copy-paste-ready next to the claims.
            for tbl in report.evidence_tables.get(rid, []):
                if not tbl.columns:
                    continue
                block.append("")  # blank line: terminate the list before the table
                if tbl.title:
                    block.append(f"**{_mdx_escape(tbl.title)}**")
                    block.append("")
                block.append("| " + " | ".join(_mdx_escape(c) for c in tbl.columns) + " |")
                block.append("|" + "---|" * len(tbl.columns))
                block.extend(
                    "| " + " | ".join(_mdx_escape(c) for c in row) + " |"
                    for row in tbl.rows
                )
                if tbl.truncated:
                    block.append(f"\n*(first {len(tbl.rows)} rows shown)*")
            blocks.append("\n".join(block))
        parts.append("\n\n".join(blocks))

    # ## EDA — reserved for future exploratory visuals; charts wrapped as MDX
    # components will render here. Emitted only when it has content; omitted today.

    if report.data_sources:
        lines = ["## Data Sources", "| source | type | detail |", "|---|---|---|"]
        for ds in report.data_sources:
            d = ds.detail
            bits = []
            if d.get("tables"):
                bits.append("tables: " + ", ".join(_mdx_escape(t) for t in d["tables"]))
            if d.get("columns"):
                bits.append(f"{len(d['columns'])} columns")
            type_label = _SOURCE_TYPE_LABELS.get(ds.source_type, ds.source_type or "—")
            lines.append(
                f"| {_mdx_escape(ds.name)} | {type_label} | {' · '.join(bits) or '—'} |"
            )
        parts.append("\n".join(lines))

    # ## Notes & Limitations — dropped from the rendered report 2026-07-09 (team
    # decision: compact report). caveats/open_questions still populate the
    # AnalysisReport JSON body; only the markdown section is gone.
    # if report.caveats or report.open_questions:
    #     lines = ["## Notes & Limitations"]
    #     for n in report.caveats:
    #         lines.append(f"- {_mdx_escape(n.text)}")
    #     for n in report.open_questions:
    #         lines.append(f"- Open: {_mdx_escape(n.text)}")
    #     parts.append("\n".join(lines))

    # ## Attempted, Unresolved — dropped from the rendered report 2026-07-09 (team
    # decision: compact report). Failed runs still populate `report.unresolved`
    # (JSON body) and the /records curation list; only the markdown section is gone.
    # if report.unresolved:
    #     lines = [
    #         "## Attempted, Unresolved",
    #         "*These analyses ran but produced no usable evidence;"
    #         " they are not reflected in the findings above.*",
    #         "",
    #     ]
    #     lines.extend(f"- {_mdx_escape(n.text)}" for n in report.unresolved)
    #     parts.append("\n".join(lines))

    if report.excluded:
        lines = [
            "## Excluded Analyses",
            "*Excluded from this report at generation time.*",
            "",
        ]
        lines.extend(f"- {_mdx_escape(n.text)}" for n in report.excluded)
        parts.append("\n".join(lines))

    # ## How This Was Analyzed — dropped from the rendered report 2026-07-09 (team
    # decision: compact report). method_steps (and _STAGE_LABELS above) stay for the
    # AnalysisReport JSON body; only the markdown section is gone.
    # if report.method_steps:
    #     lines = ["## How This Was Analyzed"]
    #     for stage_key, label in _STAGE_LABELS:
    #         steps = [s for s in report.method_steps if s.stage == stage_key]
    #         if not steps:
    #             continue
    #         rendered = "; ".join(
    #             f"{', '.join(_mdx_escape(t) for t in s.tools_used) or '—'} ({s.status})"
    #             for s in steps
    #         )
    #         lines.append(f"**{label}** — {rendered}")
    #     parts.append("\n".join(lines))

    return "\n\n---\n\n".join(parts)


# --------------------------------------------------------------------------- #
# Service
# --------------------------------------------------------------------------- #


class ReportGenerator:
    """Generates an `AnalysisReport` from persisted records. Inject deps for tests."""

    def __init__(
        self,
        record_store=None,
        structured_chain: Runnable | None = None,
        catalog_store=None,
    ) -> None:
        self._record_store = record_store
        self._chain = structured_chain
        self._catalog_store = catalog_store

    def _ensure_record_store(self):
        if self._record_store is None:
            from ..slow_path.store import PostgresReportInputStore

            self._record_store = PostgresReportInputStore()
        return self._record_store

    def _ensure_chain(self) -> Runnable:
        if self._chain is None:
            self._chain = _get_default_chain()
        return self._chain

    def _ensure_catalog_store(self):
        if self._catalog_store is None:
            from src.catalog.store import CatalogStore

            self._catalog_store = CatalogStore()
        return self._catalog_store

    async def generate(
        self,
        analysis_id: str,
        user_id: str | None = None,
        problem_statement: ProblemStatement | None = None,
        user_name: str | None = None,
        exclude_record_ids: list[str] | None = None,
    ) -> AnalysisReport:
        all_records = await self._ensure_record_store().list_for_analysis(analysis_id)
        excluded_ids = set(exclude_record_ids or [])
        excluded = [r for r in all_records if r.record_id in excluded_ids]
        kept = [r for r in all_records if r.record_id not in excluded_ids]
        # The report body reflects only substantive runs — those with a successful
        # analysis step (the same set the report floor validates). Fully-failed runs
        # can't contradict the real findings, but they are not dropped silently
        # either: they surface in the JSON `unresolved` list and the /records
        # curation endpoint (the rendered markdown section was dropped 2026-07-09).
        records = [r for r in kept if has_successful_analysis(r)]
        unresolved_records = [r for r in kept if not has_successful_analysis(r)]
        if not records:
            raise ReportError(f"no analyses recorded for {analysis_id!r} yet")

        ps = problem_statement or ProblemStatement()
        reply_language = detect_reply_language(
            None, goal_texts=[ps.objective, *ps.business_questions]
        )
        findings = _collect_findings(records)
        caveats = _collect_notes(records, "caveats", _MAX_CAVEATS)
        open_questions = _collect_notes(records, "open_questions", _MAX_OPEN_QUESTIONS)
        method_steps = _collect_method_steps(records)
        data_sources = _build_data_sources(
            records, await self._read_catalog(user_id, analysis_id)
        )
        executive_summary, bq_answers = await self._summarize(
            ps, records, caveats, reply_language
        )

        report = AnalysisReport(
            analysis_id=analysis_id,
            user_id=user_id,
            user_name=user_name,
            version=0,  # assigned by ReportStore.save under the advisory lock
            generated_at=datetime.now(UTC),
            problem_statement=ps,
            record_ids=[r.record_id for r in records],
            record_goals={r.record_id: r.goal_restated for r in records},
            executive_summary=executive_summary,
            bq_answers=bq_answers,
            findings=findings,
            caveats=caveats,
            open_questions=open_questions,
            unresolved=[_unresolved_note(r) for r in unresolved_records],
            excluded=[_excluded_note(r) for r in excluded],
            evidence_tables=_collect_evidence(records),
            data_sources=data_sources,
            method_steps=method_steps,
        )
        report.rendered_markdown = _render_markdown(report)
        logger.info(
            "report generated",
            analysis_id=analysis_id,
            records=len(records),
            findings=len(findings),
        )
        return report

    async def _read_catalog(self, user_id: str | None, analysis_id: str | None):
        """Prefer the analysis-scope catalog (this analysis's bound sources + their
        real names); fall back to the user-scope catalog when the analysis has no row
        (legacy / unbound)."""
        try:
            store = self._ensure_catalog_store()
            if analysis_id:
                cat = await store.get_by_analysis(analysis_id)
                if cat is not None:
                    return cat
            return await store.get(user_id) if user_id else None
        except Exception as exc:  # data_sources falls back; never break the report
            logger.warning("catalog read failed; data_sources will fall back", error=str(exc))
            return None

    async def _summarize(
        self,
        ps: ProblemStatement,
        records: list[AnalysisRecord],
        caveats: list[AttributedNote],
        reply_language: str,
    ) -> tuple[str, list[BusinessQuestionAnswer]]:
        human_content = _build_human_content(ps, records, caveats, reply_language)
        try:
            narrative: ReportSummaryNarrative = await self._ensure_chain().ainvoke(
                {"human_content": human_content}
            )
        except Exception as exc:  # D1: degrade, don't fail the whole report
            logger.warning("report summary LLM failed; using fallback", error=repr(exc))
            return _FALLBACK_SUMMARY, []
        return narrative.executive_summary, _resolve_bq_answers(
            narrative.bq_answers, ps.business_questions, records
        )