from __future__ import annotations
from typing import Literal, Optional, List, Dict, Any
from pydantic import BaseModel, Field

NodeStatus = Literal["ongoing", "completed"]


class NodeData(BaseModel):
    id: str
    label: str
    description: str = ""
    status: NodeStatus = "ongoing"
    depth: int = 1
    complexity: int = Field(3, ge=1, le=5, description="Conceptual density: 1=simple, 5=very complex")
    parent_id: Optional[str] = None
    children_ids: List[str] = Field(default_factory=list)
    # Which uploaded paper(s) this node draws from -> scopes its lessons/RAG. Empty = all.
    document_ids: List[str] = Field(default_factory=list)
    # Canonical page/bbox-grounded evidence behind this node. These IDs let
    # lessons, chat, graphs, and citations transfer the same source anchors.
    evidence_ids: List[str] = Field(default_factory=list)
    project_memory_ids: List[str] = Field(default_factory=list)
    # True when this node synthesizes overlapping treatments from 2+ source documents
    # (see BrainAgent.cleanup_curriculum) -> tutoring agents use this to avoid
    # conflating the papers' individual treatments when explaining the concept.
    is_merged: bool = Field(False)
    merge_summary: str = Field("", description="What's shared vs. distinct across source documents, when is_merged is True")
    # "exploration" = spawned by the GraphCuratorAgent from off-curriculum student
    # activity (see COMMIT_PROJECT) -> rendered as a distinct color from planned
    # curriculum nodes so the student can see what they discovered vs. what was planned.
    origin: Literal["curriculum", "exploration"] = "curriculum"
    # How the student's cross-session Cognee memory reshaped this node when the
    # curriculum was generated -> this is what makes memory LOAD-BEARING on the
    # graph's structure (not just soft prompt context). "new" = genuinely new
    # material; "review" = a light pass because prior sessions show it's largely
    # mastered; "scaffold" = inserted to shore up a gap memory surfaced. Brain Agent
    # sets it during generation from query_prior_knowledge (see BrainAgent).
    memory_tag: Literal["new", "review", "scaffold"] = "new"


class KnowledgeEdge(BaseModel):
    id: str
    source: str
    target: str
    edge_type: Literal["prerequisite", "related", "contains"] = "prerequisite"


class NodePatch(BaseModel):
    node_id: str
    status: Optional[NodeStatus] = None
    updated_label: Optional[str] = None
    updated_description: Optional[str] = None
    new_children: Optional[List[str]] = Field(None, description="IDs of newly generated sub-topics")


class FormulaStep(BaseModel):
    latex: str
    explanation: str


class FormulaContent(BaseModel):
    main_latex: str
    steps: List[FormulaStep] = Field(default_factory=list)


class TextRefContent(BaseModel):
    body_markdown: str
    source_label: str = ""
    source_url: str = ""


class HTML5VisualPayload(BaseModel):
    html_code: str = ""   # iframe path: graph / 3d / 2d_anim / decline
    animation_type: Literal["formula", "graph", "2d_text", "3d", "2d_anim"]
    explanation: str = Field("", description="A 2-3 sentence explanation describing exactly what the visualization demonstrates, how to use the interactive elements/sliders, and how it connects to the study material.")
    formula: Optional[FormulaContent] = None
    text_ref: Optional[TextRefContent] = None
    # Honest transparency signal for the frontend: "paper" only when the engine
    # actually verified specific content against the uploaded source (a verbatim
    # anchor match, or -- for D3 charts, which have no anchor concept -- at least
    # one generated label/value found verbatim in the source chunks). "web" is
    # reserved for when a cited external source was used instead (TextRefEngine
    # already accepts web_results; not yet wired to a live caller). Anything else
    # is illustrative/from the model's own knowledge and must say so.
    source: Literal["paper", "web", "model_knowledge"] = "model_knowledge"
    # App-owned registered templates are deterministic HTML and must never enter
    # the generated-code sandbox repair loop.
    trusted_template: bool = False
    template_id: str = ""


class ExternalAction(BaseModel):
    action_type: Literal["YOUTUBE_FETCH", "GENERATE_FLASHCARDS"]
    parameters: Dict[str, Any]


class OrchestratorAction(BaseModel):
    intent: Literal["UPDATE_GRAPH", "GENERATE_VISUAL", "STREAM_CHAT", "TOOL_CALL"]
    chat_stream_response: str = ""
    graph_patches: Optional[List[NodePatch]] = None
    visual_payload: Optional[HTML5VisualPayload] = None
    tool_execution: Optional[ExternalAction] = None