app/agents/nucleic_visual_agent.py

from __future__ import annotations

import logging

from app.agents.cerebras_client import CerebrasClient
from app.schemas.visual_lesson import NucleicVisualizationDraft, PlanCritique, VisualizationRequest
from app.services.nucleic_resolver import ResolvedNucleic


logger = logging.getLogger(__name__)


class NucleicVisualAgent:
    def __init__(self, client: CerebrasClient | None = None) -> None:
        self.client = client or CerebrasClient()

    @staticmethod
    def fallback(request: VisualizationRequest, resolved: ResolvedNucleic) -> NucleicVisualizationDraft:
        text = f"{request.prompt} {request.selection_text}".lower()
        if resolved.structure:
            primary = "structure"
        elif resolved.intent.mode == "comparison":
            primary = "comparison"
        elif resolved.nucleobase_molecule and any(
            term in text for term in ("3d", "molecule", "element", "atom", "ball-and-stick", "model")
        ):
            primary = "molecule_3d"
        elif resolved.intent.focus_base:
            primary = "chemistry"
        elif any(term in text for term in ("chemistry", "chemical", "nucleotide", "atom", "bond")):
            primary = "chemistry"
        else:
            primary = "helix_3d"
        captions = {
            "structure": "Official coordinates; rotate, zoom, and focus a verified polymer chain.",
            "comparison": "Geometry and chemistry are aligned so structural differences remain visually comparable.",
            "chemistry": "Select a base to inspect its bond pattern and connection to sugar and phosphate.",
            "helix_3d": "Backbones wind in opposite directions around paired bases; select a base for its chemistry.",
        }
        if resolved.intent.focus_base and resolved.nucleobase_molecule:
            molecule = resolved.nucleobase_molecule
            caption = (
                f"{molecule.name.title()} · {molecule.molecular_formula} · PubChem CID {molecule.pubchem_cid}. "
                + ("Select an atom to inspect its element and bonds." if primary == "molecule_3d" else "The bond graph shows the isolated nucleobase, without sugar or phosphate.")
            )
        else:
            caption = captions[primary]
        return NucleicVisualizationDraft(
            primary_view=primary,
            visible_feature_ids=[feature.feature_id for feature in resolved.features[:8]],
            caption=caption,
        )

    def _plan(self, request: VisualizationRequest, resolved: ResolvedNucleic, feedback: list[str] | None = None) -> NucleicVisualizationDraft:
        catalog = "\n".join(f"- {item.feature_id}: {item.label}" for item in resolved.features)
        context = "\n".join(
            f"- {source.origin}: {source.excerpt[:280]}"
            for source in resolved.sources
            if source.origin in {"selection", "project", "web"} and source.excerpt
        )[:2400]
        messages = [
            {
                "role": "system",
                "content": (
                    "Choose a visual composition for a nucleic-acid visualization. The visual must answer the prompt; it is not a lesson. "
                    "Choose only from the supplied feature IDs and one primary view. Use helix_3d for overall form, chemistry for bonds or nucleotides, "
                    "molecule_3d only when an official nucleobase conformer exists, comparison only for an actual comparison request, and structure only "
                    "when official macromolecular coordinates exist. The caption must be one short "
                    "sentence under 180 characters. Do not create sequences, facts, identifiers, coordinates, labels, commands, or teaching steps."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"PROMPT:\n{request.prompt}\nMODE: {resolved.intent.mode}\nMOLECULE: {resolved.intent.molecule}\n"
                    f"HAS OFFICIAL STRUCTURE: {bool(resolved.structure)}\nHAS OFFICIAL NUCLEOBASE CONFORMER: "
                    f"{bool(resolved.nucleobase_molecule)}\nFEATURES:\n{catalog}\nVISUAL CONTEXT:\n{context or '- none'}"
                    + ("\nREVISION:\n" + "\n".join(feedback) if feedback else "")
                ),
            },
        ]
        return self.client.structured_complete(
            messages,
            NucleicVisualizationDraft,
            reasoning_effort="medium",
            timeout=8.0,
            trace_label=f"nucleic-plan:{request.request_id}",
        )

    def _critic(self, request: VisualizationRequest, resolved: ResolvedNucleic, draft: NucleicVisualizationDraft) -> PlanCritique:
        messages = [
            {
                "role": "system",
                "content": (
                    "Review a visual-first nucleic-acid composition. Reject it if it does not answer the prompt, selects unavailable features, "
                    "uses structure or molecule_3d without the corresponding verified coordinates, uses comparison for a non-comparison request, "
                    "or lets prose dominate the visual."
                ),
            },
            {"role": "user", "content": f"PROMPT: {request.prompt}\nPLAN: {draft.model_dump_json()}\nAVAILABLE: {[item.feature_id for item in resolved.features]}"},
        ]
        return self.client.structured_complete(
            messages,
            PlanCritique,
            reasoning_effort="low",
            timeout=8.0,
            trace_label=f"nucleic-critic:{request.request_id}",
        )

    def plan_with_review(self, request: VisualizationRequest, resolved: ResolvedNucleic) -> tuple[NucleicVisualizationDraft, list[str]]:
        warnings: list[str] = []
        logger.info("Nucleic visual planner start request=%s", request.request_id)
        try:
            draft = self._plan(request, resolved)
        except Exception as exc:
            logger.warning("Nucleic visual planner fallback request=%s error=%s", request.request_id, exc)
            return self.fallback(request, resolved), ["Visual composition planning was unavailable; a deterministic prompt-matched composition was used."]
        try:
            logger.info("Nucleic visual critic start request=%s", request.request_id)
            critique = self._critic(request, resolved, draft)
            if not critique.approved:
                logger.info("Nucleic visual revision start request=%s", request.request_id)
                draft = self._plan(request, resolved, critique.feedback)
        except Exception as exc:
            logger.warning("Nucleic visual critic fallback request=%s error=%s", request.request_id, exc)
            warnings.append("Visual composition review was unavailable; the schema-valid composition was used.")
        available = {feature.feature_id for feature in resolved.features}
        if any(feature_id not in available for feature_id in draft.visible_feature_ids):
            draft = self.fallback(request, resolved)
            warnings.append("The proposed callouts referenced unavailable features and were replaced deterministically.")
        if resolved.features and not draft.visible_feature_ids:
            draft = self.fallback(request, resolved)
        if draft.primary_view == "structure" and not resolved.structure:
            draft = self.fallback(request, resolved)
        if draft.primary_view == "molecule_3d" and not resolved.nucleobase_molecule:
            draft = self.fallback(request, resolved)
        if draft.primary_view == "comparison" and resolved.intent.mode != "comparison":
            draft = self.fallback(request, resolved)
        required_view = self.fallback(request, resolved).primary_view
        if draft.primary_view != required_view:
            draft = draft.model_copy(update={"primary_view": required_view})
        logger.info("Nucleic visual planner done request=%s view=%s", request.request_id, draft.primary_view)
        return draft, warnings