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 ..slow_path.schemas import AnalysisRecord, TaskSummary
from .errors import ReportError
from .readiness import has_successful_analysis
from .schemas import (
    AnalysisReport,
    AttributedNote,
    DataSourceRef,
    ProblemStatement,
    ReportFinding,
    ReportSummaryNarrative,
)

logger = get_logger("report_generator")

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

# 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"),
]

# 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 _collect_notes(records: list[AnalysisRecord], field: str) -> list[AttributedNote]:
    # Caveats / open_questions are deduped by text; a merged note cites every
    # record it came from (plural record_ids).
    merged: dict[str, list[str]] = {}
    for rec in records:
        for text in getattr(rec, field):
            ids = merged.setdefault(text, [])
            if rec.record_id not in ids:
                ids.append(rec.record_id)
    return [AttributedNote(text=text, record_ids=ids) for text, ids in merged.items()]


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, findings: list[ReportFinding], caveats: list[AttributedNote]
) -> str:
    sections = []
    if ps.objective:
        sections.append("# Objective\n" + ps.objective)
    if ps.business_questions:
        sections.append(
            "# Business questions\n" + "\n".join(f"- {q}" for q in ps.business_questions)
        )
    sections.append(
        "# Findings (already finalized — synthesize, do not add numbers)\n"
        + "\n".join(f"- {f.text}" for f in findings)
    )
    if caveats:
        sections.append("# Caveats\n" + "\n".join(f"- {c.text}" for c in caveats))
    return "\n\n".join(sections)


# 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.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))
            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))

    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))

    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,
    ) -> AnalysisReport:
        records = await self._ensure_record_store().list_for_analysis(analysis_id)
        # The report reflects only substantive runs — those with a successful
        # analysis step (the same set the report floor validates). Fully-failed runs
        # are dropped so their failure narration can't contradict the real findings.
        records = [r for r in records if has_successful_analysis(r)]
        if not records:
            raise ReportError(f"no analyses recorded for {analysis_id!r} yet")

        ps = problem_statement or ProblemStatement()
        findings = _collect_findings(records)
        caveats = _collect_notes(records, "caveats")
        open_questions = _collect_notes(records, "open_questions")
        method_steps = _collect_method_steps(records)
        data_sources = _build_data_sources(
            records, await self._read_catalog(user_id, analysis_id)
        )
        executive_summary = await self._summarize(ps, findings, caveats)

        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,
            findings=findings,
            caveats=caveats,
            open_questions=open_questions,
            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, findings: list[ReportFinding], caveats: list[AttributedNote]
    ) -> str:
        human_content = _build_human_content(ps, findings, caveats)
        try:
            narrative: ReportSummaryNarrative = await self._ensure_chain().ainvoke(
                {"human_content": human_content}
            )
            return narrative.executive_summary
        except Exception as exc:  # D1: degrade, don't fail the whole report
            logger.warning("report summary LLM failed; using fallback", error=str(exc))
            return _FALLBACK_SUMMARY