shared/schemas/math_agent.py

"""Cross-service contract for the Math Agent integration layer."""
from __future__ import annotations

from datetime import datetime
from typing import Any, Literal
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field

from shared.schemas.render_job import RenderQuality


class GeometryObject(BaseModel):
    """A single geometric entity the Math Agent wants visualized."""

    model_config = ConfigDict(extra="forbid")

    type: str = Field(description="Geometry type: circle, triangle, line, angle, point, polygon, function_plot, etc.")
    label: str | None = Field(default=None, description="Optional identifier or label: A, B, C, C_1, etc.")
    properties: dict[str, Any] = Field(
        default_factory=dict,
        description="Coordinates, dimensions, colors, or constraints (e.g. {'center': [0, 0], 'radius': 3})",
    )


class AnimationDirective(BaseModel):
    """One animation beat requested by the Math Agent."""

    model_config = ConfigDict(extra="forbid")

    action: str = Field(description="Animation action: draw, highlight, transform, fade_in, fade_out, write, etc.")
    targets: list[str] = Field(default_factory=list, description="Labels or names of geometry objects involved")
    narration: str | None = Field(default=None, max_length=2_000, description="Voiceover/explanation for this beat")
    duration_hint: float | None = Field(default=None, ge=0.1, le=60.0, description="Estimated duration in seconds")


class OutputConfig(BaseModel):
    """Rendering and output configuration for the generated animation."""

    model_config = ConfigDict(extra="forbid")

    quality: RenderQuality = "720p"
    format: Literal["mp4", "gif"] = "mp4"
    language: str = Field(default="vi", min_length=2, max_length=16, description="Language code: vi, en, ...")


class VisualizationSpec(BaseModel):
    """The standardized schema that Math Agent sends to the Manim Module."""

    model_config = ConfigDict(extra="forbid")

    problem: str = Field(
        min_length=1,
        max_length=20_000,
        description="Math problem description, theorem statement, or topic to visualize",
    )
    solution_steps: list[str] = Field(
        min_length=1,
        max_length=50,
        description="Ordered list of reasoning or solution steps",
    )
    geometry: list[GeometryObject] = Field(
        default_factory=list,
        description="List of geometric or mathematical entities to construct",
    )
    animations: list[AnimationDirective] = Field(
        default_factory=list,
        description="Optional animation script and narration beats",
    )
    output_config: OutputConfig = Field(default_factory=OutputConfig)

    def to_prompt(self) -> str:
        """Serialize the spec into a rich natural language prompt for the AI pipeline."""
        parts = [f"Chủ đề / Đề bài: {self.problem}"]
        if self.solution_steps:
            parts.append("Các bước giải thích / chứng minh chi tiết:")
            for idx, step in enumerate(self.solution_steps, 1):
                parts.append(f"  {idx}. {step}")
        if self.geometry:
            parts.append("Các đối tượng hình học / toán học cần trực quan hóa:")
            for obj in self.geometry:
                label_str = f" ({obj.label})" if obj.label else ""
                props_str = f" - thuộc tính: {obj.properties}" if obj.properties else ""
                parts.append(f"  - {obj.type}{label_str}{props_str}")
        if self.animations:
            parts.append("Chỉ dẫn hoạt họa và thuyết minh (animation beats):")
            for idx, anim in enumerate(self.animations, 1):
                targets_str = f" [đối tượng: {', '.join(anim.targets)}]" if anim.targets else ""
                narr_str = f" | Lời thoại: '{anim.narration}'" if anim.narration else ""
                parts.append(f"  - Beat {idx}: {anim.action}{targets_str}{narr_str}")
        return "\n".join(parts)


class MathRenderRequest(BaseModel):
    """Request payload for POST /v1/math/generate."""

    model_config = ConfigDict(extra="forbid")

    spec: VisualizationSpec
    callback_url: str | None = Field(
        default=None,
        max_length=2_048,
        description="Optional webhook URL to receive notification when video generation finishes",
    )


class MathRenderResponse(BaseModel):
    """Status or completion response for Math Agent."""

    model_config = ConfigDict(extra="forbid")

    job_id: UUID = Field(description="Identifier for tracking the generation & rendering job")
    project_id: UUID = Field(description="Corresponding Manim project ID")
    status: Literal["queued", "generating", "rendering", "completed", "failed"]
    video_url: str | None = Field(default=None, description="Signed or direct URL to the final rendered video")
    duration: float | None = Field(default=None, description="Total video duration in seconds")
    error: str | None = Field(default=None, description="Error message if generation or rendering failed")
    created_at: datetime = Field(default_factory=datetime.now)
    completed_at: datetime | None = None