app/agents/tutor_agent.py

"""Tutor Agent -> grounded lessons, sandbox repair.

Every fact in grounded_truth must cite its RAG source. Visual generation now
lives in dedicated engines (FormulaEngine, TextRefEngine, ShellVisualEngine,
D3Engine) behind VisualModalityRouter, not in this class.
"""
import re
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from app.agents.cerebras_client import CerebrasClient
from app.schemas.graph import HTML5VisualPayload, NodeData


class LessonPayload(BaseModel):
    anchor: str = Field(description="Conceptual introduction tailored to familiarity level")
    grounded_truth: str = Field(
        description="Facts with inline [Source: X, chunk N] citations from the RAG chunks"
    )
    citations: List[str]
    visual_suggestion: str = Field(
        description="One of: three.js, canvas, katex, plot, quote, none"
    )


class _Flashcard(BaseModel):
    front: str
    back: str
    source_chunk_indexes: List[int] = Field(description="Indexes of the chunks used to synthesize this flashcard")


class _FlashcardsPayload(BaseModel):
    cards: List[_Flashcard]


_BANNED_FLASHCARD_PHRASES = (
    "the text explicitly states",
    "according to the chunk",
    "according to the source",
    "the source mentions",
    "the chunk mentions",
    "is defined as",
)


def _normalized_flashcard_front(text: str) -> str:
    return " ".join(text.split()).casefold()


def _flashcard_rejection_reason(
    card: _Flashcard,
    chunk_count: int,
    seen_fronts: set[str],
) -> str | None:
    front = card.front.strip()
    back = card.back.strip()
    if not front or not back:
        return "front and back must both be non-empty"

    normalized_front = _normalized_flashcard_front(front)
    if normalized_front in seen_fronts:
        return "question duplicates an accepted card"

    combined = f"{front}\n{back}"
    if front.count("$") % 2 or back.count("$") % 2:
        return "math delimiters are unbalanced"
    if re.search(r"\\v(?![A-Za-z])", combined):
        return r"invalid \v LaTeX command"

    lowered = combined.casefold()
    if any(phrase in lowered for phrase in _BANNED_FLASHCARD_PHRASES):
        return "card refers to the source instead of stating the concept directly"

    indexes = card.source_chunk_indexes
    if chunk_count == 0 and indexes:
        return "source indexes must be empty when no chunks are available"
    if any(type(index) is not int or index < 0 or index >= chunk_count for index in indexes):
        return "source index is outside the supplied chunk range"
    return None


class _MCQOption(BaseModel):
    text: str
    is_correct: bool


class _MCQ(BaseModel):
    question: str
    options: List[_MCQOption]
    explanation: str
    source_chunk_indexes: List[int] = Field(description="Indexes of the chunks used to synthesize this question")


class _QuizPayload(BaseModel):
    questions: List[_MCQ]


class _MCQEvaluation(BaseModel):
    is_good: bool = Field(description="True if question is high-quality, conceptual, and stands alone.")
    has_latex_errors: bool = Field(description="True if there are malformed LaTeX commands or mismatched delimiters.")
    has_hallucinated_terms: bool = Field(description="True if it uses terms not present in the source.")
    reason: str


class _QuizQualityPayload(BaseModel):
    evaluations: List[_MCQEvaluation]


class TutorAgent:
    def __init__(self, client: Optional[CerebrasClient] = None) -> None:
        self._client = client or CerebrasClient()

    # ------------------------------------------------------------------ #
    # Lesson                                                               #
    # ------------------------------------------------------------------ #

    async def stream_lesson(
        self,
        node: NodeData,
        chunks: List[Dict[str, Any]],
        familiarity: str,
        knowledge_mode: str = "content_only",
        web_context: str = "",
        prior_knowledge: str = "",
    ):
        """Async generator that yields lesson text tokens for streaming to the client.

        prior_knowledge is this student's cross-session Cognee memory for THIS concept
        (from query_prior_knowledge). When present, the lesson adapts per-concept ->
        it doesn't re-explain what they've already shown they grasp and leans into the
        parts they've struggled with. Tier 1A load-bearing tenet: the Tutor adapts from
        cross-document history, not just the once-per-session familiarity setting.
        """
        chunk_text = "\n\n".join(
            f"[Source: {c['source']}, chunk {c.get('chunk_index', i)}]\n{c['text']}"
            for i, c in enumerate(chunks)
        )
        # Familiarity-aware math formatting instruction
        if familiarity in ("graduate", "expert"):
            math_note = (
                "You may use LaTeX math notation with $...$ for inline math and $$...$$ for display math. "
                "Use proper LaTeX commands (e.g. \\frac{a}{b}, \\sqrt{x}, \\hat{m}_t)."
            )
        else:
            math_note = (
                "Describe all mathematical concepts in plain words and simple notation. "
                "Do NOT use LaTeX $ delimiters or complex formulas. "
                "For example, say 'a divided by b' instead of $\\frac{a}{b}$."
            )

        if knowledge_mode == "net_support":
            grounding = (
                "Ground your explanation in the provided SOURCE MATERIAL and WEB SOURCE MATERIAL only. "
                "If the topic is missing from the student's uploaded content, draw on the web search results "
                "to fill the gap -> cite the web source, not your own training weights. "
                "Do NOT invent facts not present in either the source chunks or the web results."
            )
        else:
            grounding = (
                "Base the lesson EXCLUSIVELY on the provided source material (the chunks below). "
                "Do NOT generate facts or claims from your own training weights. "
                "If the source material does not directly cover a point, say so and "
                "anchor your explanation to the closest relevant passage that IS in the source."
            )
        merge_note = ""
        if node.is_merged:
            merge_note = (
                f"\n\nNOTE: '{node.label}' synthesizes overlapping treatments from multiple source "
                f"documents. {node.merge_summary} When explaining, be explicit about which document a "
                "specific claim, method, or result comes from if the source chunks reflect differing "
                "treatments (each chunk below is tagged with its source) -> do not blend the documents' "
                "framings into one as if they were a single source."
            )
        memory_note = ""
        if prior_knowledge:
            memory_note = (
                "\n\nWHAT THIS STUDENT ALREADY KNOWS ABOUT THIS CONCEPT (from prior sessions -> adapt, "
                "don't re-teach it): skim what they clearly have, and spend the lesson on the parts they "
                "haven't engaged or have struggled with. Do NOT announce that you're adapting; just teach "
                f"at the right level.\n{prior_knowledge}"
            )
        messages = [
            {
                "role": "system",
                "content": (
                    f"You are a Cognitive Translator and tutor. {grounding}{merge_note}{memory_note}\n\n"
                    "Write a single, flowing lesson -> do NOT split into separate sections like "
                    "'Concept' or 'From the Source'. Weave the intuitive explanation and factual "
                    "details together naturally into one coherent narrative.\n\n"
                    "FORMATTING RULES:\n"
                    f"- {math_note}\n"
                    "- Use **bold** for key terms.\n"
                    "- Use bullet points (lines starting with '* ') for lists.\n"
                    "- Do NOT include [Source: X, chunk N] citations or Web source URLs in your output -> "
                    "the student should not see internal source or web reference links."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Teach '{node.label}' at {familiarity} level.\n\n"
                    f"SOURCE MATERIAL:\n{chunk_text}"
                ),
            },
        ]
        if web_context:
            messages[-1]["content"] += f"\n\nWEB SOURCE MATERIAL:\n{web_context}"

        async for token in self._client.stream_complete(messages):
            yield token

    def generate_lesson(
        self,
        node: NodeData,
        chunks: List[Dict[str, Any]],
        familiarity: str,
    ) -> LessonPayload:
        chunk_text = "\n\n".join(
            f"[Source: {c['source']}, chunk {c.get('chunk_index', i)}]\n{c['text']}"
            for i, c in enumerate(chunks)
        )
        messages = [
            {
                "role": "system",
                "content": (
                    "You are a Cognitive Translator. Use ONLY the provided source material. "
                    "Every fact in grounded_truth MUST have an inline citation like "
                    "[Source: X, chunk N]. Never hallucinate. If the source doesn't cover "
                    "something, say so explicitly."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Teach '{node.label}' at {familiarity} level.\n\n"
                    f"SOURCE MATERIAL:\n{chunk_text}\n\n"
                    "Write anchor (intuitive intro), grounded_truth (cited facts), "
                    "and pick a visual_suggestion."
                ),
            },
        ]
        return self._client.structured_complete(messages, LessonPayload)

    # ------------------------------------------------------------------ #
    # Flashcards + Quiz                                                   #
    # ------------------------------------------------------------------ #

    def generate_flashcards(
        self, node_label: str, chunks: List[Dict[str, Any]], familiarity: str
    ) -> _FlashcardsPayload:
        chunk_text = "\n\n".join(
            f"[Chunk {c.get('chunk_index', i)}]\n{c['text']}" for i, c in enumerate(chunks)
        )
        user_content: List[Dict[str, Any]] = [
            {
                "type": "text",
                "text": (
                    f"Create deep, conceptual flashcards about '{node_label}' at the {familiarity} level.\n\n"
                    f"SOURCE:\n{chunk_text}"
                ),
            }
        ]

        messages = [
            {
                "role": "system",
                "content": (
                    "Generate exactly 10 high-level, conceptual open-recall flashcards. "
                    "Do NOT focus on narrow trivia or fill-in-the-blanks. Synthesize information across the chunks. "
                    "front = conceptual question, back = comprehensive answer. "
                    "YOU ARE THE TUTOR. State facts directly as your own knowledge. "
                    "BANNED PHRASES: 'The text explicitly states', 'According to the chunk', 'The source mentions', 'is defined as'. "
                    "Instead of saying 'The text states X is Y', just say 'X is Y'. "
                    "Cite the chunks you used in source_chunk_indexes. "
                    "Format math using $...$ for inline and $$...$$ for block math. Do not use invalid commands like \\v."
                ),
            },
            {
                "role": "user",
                "content": user_content,
            },
        ]
        
        raw_payload = self._client.structured_complete(
            messages, _FlashcardsPayload, model="gemma-4-31b"
        )

        accepted: list[_Flashcard] = []
        seen_fronts: set[str] = set()
        rejection_reasons: list[str] = []

        def accept_batch(cards: List[_Flashcard]) -> None:
            for card_index, card in enumerate(cards):
                if len(accepted) >= 10:
                    return
                reason = _flashcard_rejection_reason(card, len(chunks), seen_fronts)
                if reason:
                    rejection_reasons.append(f"Card {card_index}: {reason}")
                    continue
                accepted.append(card)
                seen_fronts.add(_normalized_flashcard_front(card.front))

        accept_batch(raw_payload.cards)
        needed = 10 - len(accepted)
        if needed > 0:
            accepted_fronts = "\n".join(f"- {card.front.strip()}" for card in accepted) or "- None"
            reasons = "\n".join(f"- {reason}" for reason in rejection_reasons) or "- Too few cards were returned"
            repair_messages = [
                *messages,
                {"role": "assistant", "content": raw_payload.model_dump_json()},
                {
                    "role": "user",
                    "content": (
                        f"Generate exactly {needed} replacement flashcards. "
                        "Return only new cards that satisfy the original source-grounding rules.\n\n"
                        f"Accepted questions; do not repeat them:\n{accepted_fronts}\n\n"
                        f"Deterministic validation failures to repair:\n{reasons}"
                    ),
                },
            ]
            repaired_payload = self._client.structured_complete(
                repair_messages, _FlashcardsPayload, model="gemma-4-31b"
            )
            accept_batch(repaired_payload.cards)

        return _FlashcardsPayload(cards=accepted[:10])

    def generate_quiz(
        self, node_label: str, chunks: List[Dict[str, Any]], familiarity: str
    ) -> _QuizPayload:
        chunk_text = "\n\n".join(
            f"[Chunk {c.get('chunk_index', i)}]\n{c['text']}" for i, c in enumerate(chunks)
        )
        user_content: List[Dict[str, Any]] = [
            {
                "type": "text",
                "text": (
                    f"Create sophisticated multiple-choice questions about '{node_label}' at the {familiarity} level.\n\n"
                    f"SOURCE:\n{chunk_text}"
                ),
            }
        ]

        messages = [
            {
                "role": "system",
                "content": (
                    "Generate high-level, conceptual multiple-choice questions. "
                    "Synthesize information across chunks. Do not ask for verbatim quotes. "
                    "Each question has exactly 1 correct option and 3 distractors. "
                    "Include a conceptual explanation. YOU ARE THE TUTOR. State facts directly as your own knowledge. "
                    "BANNED PHRASES: 'The text explicitly states', 'According to the chunk', 'The source mentions'. "
                    "Explain the concept directly (e.g., 'AdaGrad corresponds to...' instead of 'The text states that AdaGrad...'). "
                    "Cite the chunks used in source_chunk_indexes. "
                    "Format math using $...$ for inline and $$...$$ for block math. Do not use invalid commands like \\v."
                ),
            },
            {
                "role": "user",
                "content": user_content,
            },
        ]
        
        good_qs = []
        for attempt in range(3):
            needed = 8 - len(good_qs)
            if needed <= 0:
                break
                
            # Temporarily instruct how many to generate
            if attempt > 0:
                messages[-1]["content"] += f"\n\n(Generate exactly {needed} new multiple-choice questions)"
            else:
                messages[-1]["content"] = user_content

            raw_payload = self._client.structured_complete(messages, _QuizPayload, model="gemma-4-31b")
            
            # Quality Check Pass
            qc_messages = [
                {
                    "role": "system",
                    "content": (
                        "Evaluate each quiz question strictly. "
                        "1. Does it make logical sense independently? "
                        "2. Are there invalid LaTeX commands or mismatched $? "
                        "3. Does it hallucinate technical or biological terms? "
                        "4. Are distractors plausible but clearly wrong? "
                        "5. Does the explanation use banned phrases like 'The text states', 'According to', 'As mentioned'? "
                        "Reject low-quality trivia, broken LaTeX, references to the source text in the explanation, or hallucinations."
                    )
                },
                {
                    "role": "user",
                    "content": "Questions:\n" + "\n".join(f"{i}. Q: {q.question}\nA: {q.options}" for i, q in enumerate(raw_payload.questions))
                }
            ]
            qc_payload = self._client.structured_complete(qc_messages, _QuizQualityPayload, model="gemma-4-31b")
            
            feedback = []
            batch_good = []
            for idx, (q, eval_res) in enumerate(zip(raw_payload.questions, qc_payload.evaluations)):
                if eval_res.is_good and not eval_res.has_latex_errors and not eval_res.has_hallucinated_terms:
                    batch_good.append(q)
                else:
                    feedback.append(f"Q{idx}: Rejected. Reason: {eval_res.reason}. LaTeX Error: {eval_res.has_latex_errors}. Hallucination: {eval_res.has_hallucinated_terms}")
            
            good_qs.extend(batch_good)
                    
            if len(good_qs) < 8 and attempt < 2:
                messages.append({"role": "assistant", "content": raw_payload.model_dump_json()})
                messages.append({"role": "user", "content": (
                    f"Out of that batch, {len(batch_good)} were accepted. The following were rejected by the QC Agent:\n"
                    + "\n".join(feedback) + "\n\n"
                    f"Please generate {8 - len(good_qs)} NEW questions. Fix the LaTeX and terminology errors mentioned above. Do not repeat accepted questions."
                )})

        # Fallback if too strict
        if not good_qs:
            good_qs = raw_payload.questions[:5]

        return _QuizPayload(questions=good_qs[:8])

    # ------------------------------------------------------------------ #
    # Visuals                                                             #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _decline_html(concept: str, reason: str) -> str:
        """A calm, on-brand 'no figure for this one' card -> never a red error, never a
        fabricated diagram. Self-contained, matches the app's cream/navy styling."""
        safe_concept = concept.replace("<", "&lt;").replace(">", "&gt;")
        safe_reason = reason.replace("<", "&lt;").replace(">", "&gt;")
        return (
            "<!DOCTYPE html><html><head><meta charset='utf-8'>"
            "<style>html,body{margin:0;height:100%}"
            "body{display:flex;align-items:center;justify-content:center;background:#0f0f0f;"
            "font-family:Georgia,'Libre Caslon Text',serif;color:#e2e8f0;padding:24px;box-sizing:border-box}"
            ".card{max-width:420px;text-align:center;line-height:1.6}"
            ".t{font-size:20px;color:#8fb4de;margin-bottom:12px}"
            ".r{font-size:15px;color:#cbd5e1}</style></head><body>"
            f"<div class='card'><div class='t'>{safe_concept}</div>"
            f"<div class='r'>{safe_reason}</div></div></body></html>"
        )

    def repair_visual(
        self, original_html: str, error_message: str
    ) -> HTML5VisualPayload:
        system_prompt = (
            "You are debugging a self-contained HTML5 visualisation. "
            "Fix the JavaScript error and return the corrected complete HTML."
        )
        if error_message.startswith("BlankRender:"):
            system_prompt = (
                "You are debugging a self-contained HTML5 visualisation. The code did NOT crash "
                "-- it ran successfully but produced no visible output, so there is no stack trace "
                "or line number to chase. This is NOT a syntax/runtime bug to fix in the "
                "traditional sense: it means the content-generating code likely built something "
                "but never attached it to the target the shell actually renders (for example, the "
                "model forgot to call `group.add(...)`, or a 2D animation's `draw()` was defined "
                "but never actually draws anything).\n\n"
                "Re-examine the original code specifically for:\n"
                "(a) content-building statements that create objects (meshes/lines/points, or "
                "drawing calls) but never call `.add(...)` on the right target;\n"
                "(b) a missing or malformed `return { ... }` statement -- for a 2D animation shell "
                "the code MUST return `{ draw(ctx, width, height, time, dt) { ... } }` and that "
                "draw function must actually draw something;\n"
                "(c) content added to the wrong variable, e.g. `scene.add(...)` instead of "
                "`group.add(...)` for a 3D shell, if that distinction is visible in the code.\n\n"
                "Fix the code so it actually attaches/renders visible content, then return the "
                "corrected complete HTML."
            )
        messages = [
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": (
                    f"Error: {error_message}\n\n"
                    f"Original code:\n{original_html}\n\nFix the error."
                ),
            },
        ]
        return self._client.structured_complete(messages, HTML5VisualPayload)