app/agents/protein_lesson_agent.py

from __future__ import annotations

from app.agents.cerebras_client import CerebrasClient
from app.schemas.visual_lesson import PlanCritique, ProteinLessonDraft, VisualizationRequest
from app.services.protein_resolver import ResolvedProtein


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

    @staticmethod
    def fallback(resolved: ResolvedProtein) -> ProteinLessonDraft:
        structure = resolved.structure
        return ProteinLessonDraft(
            title=f"Structure of {resolved.identity.protein_name}",
            interpretation=(
                f"Explore {structure.structure_id} from the overall fold down to its verified chains, "
                "ligands, and requested residue features."
            ),
            assumptions=[
                "Only features verified against official structure metadata are interactive.",
                "A molecular structure is a model of the archived experiment or prediction, not a live molecule.",
            ],
            teaching_steps=[
                "Confirm the protein identity and organism before interpreting the structure.",
                "Read the overall fold and how secondary-structure elements organize the molecule.",
                "Separate the polymer chains and any mapped structural regions.",
                "Inspect verified ligands, binding features, and requested residues.",
                "Return to the evidence ledger to distinguish database facts from interpretation.",
            ],
            feature_ids=[feature.feature_id for feature in resolved.features[:12]],
            initial_representation="cartoon",
            initial_color_scheme="confidence" if structure.confidence_mode == "plddt" else "chain",
        )

    @staticmethod
    def _catalog(resolved: ResolvedProtein) -> str:
        features = "\n".join(
            f"- {item.feature_id}: {item.kind}; {item.label}; mapping={item.mapping_note or 'database selector'}"
            for item in resolved.features
        )
        claims = "\n".join(f"- {claim.claim_id}: {claim.text} [{claim.support_level}]" for claim in resolved.claims)
        sources = "\n".join(
            f"- {source.source_id} ({source.origin}, {source.title}): {source.excerpt[:500]}"
            for source in resolved.sources[:16]
        )
        return (
            f"IDENTITY: {resolved.identity.model_dump_json()}\n"
            f"STRUCTURE: {resolved.structure.model_dump_json()}\n"
            f"VERIFIED FEATURES:\n{features or '- none'}\n"
            f"EVIDENCE CLAIMS:\n{claims}\n"
            f"RETRIEVED SOURCE CONTEXT:\n{sources or '- none'}"
        )[:14000]

    def _plan(self, request: VisualizationRequest, resolved: ResolvedProtein, profile: str, feedback: list[str] | None = None) -> ProteinLessonDraft:
        messages = [
            {
                "role": "system",
                "content": (
                    "Plan a compact Protein Laboratory lesson using only the supplied identity, structure, feature IDs, "
                    "and evidence claims. Never create identifiers, residues, chains, coordinates, selectors, Mol* commands, "
                    "or biological facts. Return exactly five teaching steps. Make experimental-versus-predicted status explicit. "
                    "Choose confidence coloring only when confidence_mode is plddt."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"REQUEST:\n{request.prompt}\nFAMILIARITY: {request.familiarity}\nPROFILE:\n{profile[:1400]}\n\n"
                    f"{self._catalog(resolved)}"
                    + (f"\n\nREVISION FEEDBACK:\n" + "\n".join(f"- {item}" for item in feedback) if feedback else "")
                ),
            },
        ]
        return self.client.structured_complete(messages, ProteinLessonDraft, reasoning_effort="high")

    def _critic(self, draft: ProteinLessonDraft, resolved: ResolvedProtein) -> PlanCritique:
        messages = [
            {
                "role": "system",
                "content": (
                    "Review this Protein Laboratory plan. Approve only if it has five concise steps, uses only catalogued "
                    "feature IDs, discloses predicted versus experimental status, avoids unsupported molecular claims, and "
                    "fits a narrow research side panel."
                ),
            },
            {"role": "user", "content": f"PLAN:\n{draft.model_dump_json()}\n\n{self._catalog(resolved)}"},
        ]
        return self.client.structured_complete(messages, PlanCritique, reasoning_effort="medium")

    def plan_with_review(
        self, request: VisualizationRequest, resolved: ResolvedProtein, profile: str = ""
    ) -> tuple[ProteinLessonDraft, list[str]]:
        warnings: list[str] = []
        try:
            draft = self._plan(request, resolved, profile)
        except Exception:
            return self.fallback(resolved), ["The protein planner was unavailable; a deterministic structure lesson was used."]
        try:
            critique = self._critic(draft, resolved)
            if not critique.approved:
                draft = self._plan(request, resolved, profile, critique.feedback)
        except Exception:
            warnings.append("Protein plan review was unavailable; the schema-valid plan was used.")
        return draft, warnings