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

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

_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]:
    # Findings are distinct insights — not deduped; each traces to its record.
    return [
        ReportFinding(text=text, record_ids=[rec.record_id])
        for rec in records
        for text in rec.findings
    ]


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, bound_ids: list[str] | None = None
) -> list[DataSourceRef]:
    """Freeze real catalog metadata for the sources this analysis used.

    When the analysis has a data-source binding (#10), the candidate set is scoped
    to the bound sources first (fail-open if the binding doesn't intersect the
    catalog). Within that set, matches catalog sources against the records'
    (narrative) `data_used` by name/id; falls back to all (bound) 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
    if bound_ids:
        scoped = [s for s in candidates if s.source_id in set(bound_ids)]
        candidates = scoped or candidates  # fail-open if binding doesn't match catalog

    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 = []
    ps_lines = [v for v in (ps.objective, ps.target_value, ps.scope) if v]
    if ps_lines:
        sections.append("# Problem Statement\n" + "\n".join(ps_lines))
    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)


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.
    parts: list[str] = ["# Analysis Report"]
    parts.append(
        f"*Generated {report.generated_at:%Y-%m-%d} · "
        f"{len(report.record_ids)} analyses · {len(report.data_sources)} source(s)*"
    )

    ps = report.problem_statement
    ps_lines = [v for v in (ps.objective, ps.target_value, ps.scope) if v]
    if ps_lines:
        parts.append("## Problem Statement\n" + " ".join(ps_lines))

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

    if report.findings:
        lines = ["## Key Findings"]
        for i, f in enumerate(report.findings, 1):
            cite = f" *({', '.join(f.record_ids)})*" if f.record_ids else ""
            lines.append(f"{i}. {f.text}{cite}")
        parts.append("\n".join(lines))

    if report.caveats or report.open_questions:
        lines = ["## Caveats & Open Questions"]
        for n in report.caveats:
            cite = f" *({', '.join(n.record_ids)})*" if n.record_ids else ""
            lines.append(f"- {n.text}{cite}")
        for n in report.open_questions:
            cite = f" *({', '.join(n.record_ids)})*" if n.record_ids else ""
            lines.append(f"- Open: {n.text}{cite}")
        parts.append("\n".join(lines))

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

    if report.method_steps:
        lines = ["## Appendix B — Method"]
        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(s.tools_used) or '—'} ({s.status})" for s in steps
            )
            lines.append(f"**{label}** — {rendered}")
        parts.append("\n".join(lines))

    return "\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,
        binding_store=None,
    ) -> None:
        self._record_store = record_store
        self._chain = structured_chain
        self._catalog_store = catalog_store
        self._binding_store = binding_store

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

            self._record_store = PostgresAnalysisStore()
        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,
    ) -> AnalysisReport:
        records = await self._ensure_record_store().list_for_analysis(analysis_id)
        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)
        bound_ids = await self._read_binding(analysis_id)
        data_sources = _build_data_sources(
            records, await self._read_catalog(user_id), bound_ids
        )
        executive_summary = await self._summarize(ps, findings, caveats)

        report = AnalysisReport(
            analysis_id=analysis_id,
            user_id=user_id,
            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],
            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):
        if not user_id:
            return None
        try:
            return await self._ensure_catalog_store().get(user_id)
        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

    def _ensure_binding_store(self):
        if self._binding_store is None:
            from ..binding_store import AnalysisDataSourceStore

            self._binding_store = AnalysisDataSourceStore()
        return self._binding_store

    async def _read_binding(self, analysis_id: str) -> list[str]:
        """Bound source ids for the analysis (#10). Never-throw → [] (unscoped)."""
        try:
            return await self._ensure_binding_store().get(analysis_id)
        except Exception as exc:  # data_sources falls back to whole catalog
            logger.warning("binding read failed; data_sources unscoped", error=str(exc))
            return []

    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