app/agents/d3/engine.py

from __future__ import annotations
import logging
import time
from typing import Any, Iterator, List
from pydantic import BaseModel
from app.schemas.graph import HTML5VisualPayload
from app.agents.d3 import templates as _templates  # noqa: F401 - populate trusted registry
from app.agents.d3.registry import TEMPLATES
from app.agents.d3.router import D3TemplateRouter
from app.agents.d3.extractor import D3DataExtractor

logger = logging.getLogger(__name__)


class D3MetricClarificationError(ValueError):
    def __init__(self, metrics: list[str]) -> None:
        self.metrics = metrics[:5]
        super().__init__("The selected source contains several equally plausible metrics.")


def _iter_string_leaves(obj: Any) -> Iterator[str]:
    if isinstance(obj, BaseModel):
        for v in obj.__dict__.values():
            yield from _iter_string_leaves(v)
    elif isinstance(obj, (list, tuple)):
        for v in obj:
            yield from _iter_string_leaves(v)
    elif isinstance(obj, str):
        yield obj


def _looks_grounded(data: BaseModel, chunk_text: str) -> bool:
    """D3 chart data has no anchor concept (unlike formula/text_ref/shell, which
    self-check a verbatim source substring) -- this is the closest honest signal:
    true only if at least one distinctive label the model produced actually
    appears verbatim in the source chunks.

    Two deliberate filters against false positives (a fabricated chart getting
    mislabeled as source-grounded by coincidence):
    - Short strings (< 6 chars) are excluded -- too likely to appear by chance.
    - Purely numeric strings (years, counts: "2010", "1990") are excluded even if
      long enough -- these are exactly the kind of value a synthesized chart is
      most likely to share with unrelated real content by pure coincidence (a
      fabricated "2010" data point "matching" an unrelated year mentioned
      elsewhere in the source proves nothing about whether the chart itself is
      grounded).
    """
    if not chunk_text:
        return False
    for value in _iter_string_leaves(data):
        candidate = value.strip()
        if len(candidate) < 6:
            continue
        if candidate.replace(".", "", 1).replace("-", "", 1).isdigit():
            continue
        if candidate in chunk_text:
            return True
    return False


class D3Engine:
    def __init__(self, router=None, extractor=None) -> None:
        self._router = router or D3TemplateRouter()
        self._extractor = extractor or D3DataExtractor()

    def generate(self, concept: str, chunks: List[dict], familiarity: str) -> HTML5VisualPayload:
        sel = self._router.select(concept, chunks, familiarity)
        template = TEMPLATES[sel.template_id]
        data = self._extractor.fill(template, concept, chunks, familiarity)
        html = template.render(data)
        chunk_text = "\n\n".join(c.get("text", "") for c in chunks)
        source = "paper" if _looks_grounded(data, chunk_text) else "model_knowledge"
        return HTML5VisualPayload(
            html_code=html, animation_type="graph", explanation=f"{template.title}: {concept}.", source=source
        )

    def transform_selected(
        self,
        concept: str,
        chunks: List[dict],
        familiarity: str,
        image_base64: str = "",
    ) -> tuple[HTML5VisualPayload, str]:
        started = time.perf_counter()
        selection = self._router.select_transform(concept, chunks, familiarity)
        if selection.requires_metric_clarification and len(selection.metric_candidates) >= 2:
            raise D3MetricClarificationError(selection.metric_candidates)
        routed_ms = (time.perf_counter() - started) * 1000
        template = TEMPLATES[selection.template_id]
        data = self._extractor.fill(
            template,
            concept,
            chunks,
            familiarity,
            strict_source=True,
            image_base64=image_base64,
            selected_metric=selection.selected_metric,
        )
        extracted_ms = (time.perf_counter() - started) * 1000 - routed_ms
        narrative_started = time.perf_counter()
        try:
            answer_markdown = self._extractor.explain(
                concept=concept,
                template=template,
                data=data,
                chunks=chunks,
                image_base64=image_base64,
            )
        except Exception:
            logger.exception("D3 selected transform narrative failed for template=%s", selection.template_id)
            answer_markdown = (
                f"The selected source is shown as a **{template.title.lower()}**"
                + (f" using **{selection.selected_metric}**." if selection.selected_metric else ".")
            )
        narrative_ms = (time.perf_counter() - narrative_started) * 1000
        logger.info(
            "D3 selected transform template=%s metric=%r source=%s route_ms=%.1f extract_ms=%.1f narrative_ms=%.1f total_ms=%.1f",
            selection.template_id,
            selection.selected_metric,
            "image" if image_base64 else "text",
            routed_ms,
            extracted_ms,
            narrative_ms,
            (time.perf_counter() - started) * 1000,
        )
        metric_suffix = f" using {selection.selected_metric}" if selection.selected_metric else ""
        visual = HTML5VisualPayload(
            html_code=template.render(data),
            animation_type="graph",
            explanation=f"{template.title}{metric_suffix}, transcribed from the selected source.",
            source="paper",
            trusted_template=True,
            template_id=selection.template_id,
        )
        return visual, answer_markdown