backend/app/api/v1/math_agent.py

"""Math Agent integration endpoints.

Exposes a simplified Video Generation Module interface for external Math Agents:
- POST /v1/math/generate: submit a VisualizationSpec to queue generation & rendering
- GET  /v1/math/jobs/{job_id}: poll status and retrieve rendered video URL
"""
from __future__ import annotations

import hmac
import logging
from uuid import UUID, uuid4

from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from shared.schemas.math_agent import MathRenderRequest, MathRenderResponse
from shared.schemas.render_job import RenderQuality

from app.api.deps import (
    ContentStore,
    get_content_store,
    get_hitl_store,
    get_job_store,
    get_request_user_id,
)
from app.api.v1.hitl import get_pipeline_service
from app.core.config import settings
from app.services.ai_queue import AiQueue, AiQueueUnavailable
from app.services.events import publish_project_event
from app.services.hitl_service import HitlPipelineService
from app.services.hitl_store import SupabaseHitlStore
from app.services.job_store import RedisRenderJobStore
from app.services.render_snapshot import project_render_source
from app.services.supabase_storage_rest import sign_storage_object_read_url

router = APIRouter(tags=["math-agent"])
logger = logging.getLogger(__name__)


def resolve_math_caller_user_id(
    request: Request,
    x_internal_token: str | None = Header(None, alias="X-Internal-Token"),
) -> UUID:
    """Authenticate either via internal service token or standard user session."""
    system_user = UUID("ca86b9cd-19df-4866-8c72-51223077d746")
    if x_internal_token and hmac.compare_digest(
        x_internal_token.strip(), settings.internal_service_token
    ):
        return system_user if str(settings.dev_default_user_id).startswith("00000000") else settings.dev_default_user_id
    if settings.auth_mode == "off":
        return system_user if str(settings.dev_default_user_id).startswith("00000000") else settings.dev_default_user_id
    return get_request_user_id(request)


def _resolve_video_url(project_id: UUID, raw_video_ref: str | None) -> str | None:
    """Convert an internal video reference (supabase:// or file://) to a readable URL."""
    if not raw_video_ref:
        return None
    expected_prefix = f"supabase://{settings.supabase_storage_bucket.strip()}/"
    if raw_video_ref.startswith(expected_prefix):
        object_path = raw_video_ref.removeprefix(expected_prefix).lstrip("/")
        try:
            return sign_storage_object_read_url(object_path=object_path)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Unable to sign video URL for project %s: %s", project_id, exc)
            return None
    if raw_video_ref.startswith("file://"):
        return f"/v1/projects/{project_id}/rendered-video"
    return raw_video_ref


@router.post(
    "/generate",
    response_model=MathRenderResponse,
    status_code=status.HTTP_202_ACCEPTED,
    summary="Generate a math animation video from a VisualizationSpec",
)
def generate_math_video(
    body: MathRenderRequest,
    user_id: UUID = Depends(resolve_math_caller_user_id),  # noqa: B008
    content: ContentStore = Depends(get_content_store),  # noqa: B008
    pipeline_service: HitlPipelineService = Depends(get_pipeline_service),  # noqa: B008
) -> MathRenderResponse:
    """Accepts a VisualizationSpec from a Math Agent, creates a project,

    and kicks off the automated (HITL-off) Manim video generation pipeline.
    """
    spec = body.spec
    project_id = uuid4()
    title = spec.problem[:100] if len(spec.problem) > 100 else spec.problem

    project_config = {
        "math_spec": spec.model_dump(mode="json"),
        "callback_url": body.callback_url,
        "quality": spec.output_config.quality,
    }

    project = content.create_project(
        project_id=project_id,
        user_id=user_id,
        title=title,
        description=spec.problem,
        source_language=spec.output_config.language,
        target_scenes=max(1, len(spec.solution_steps)),
        status="draft",
        config=project_config,
    )

    prompt = spec.to_prompt()
    run, _first_step = pipeline_service.start_project_run(
        project_id=project.id,
        user_id=user_id,
        prompt=prompt,
        hitl_enabled=False,
    )

    logger.info(
        "Math Agent generation started project_id=%s run_id=%s",
        project.id,
        run.id,
    )

    return MathRenderResponse(
        job_id=run.id,
        project_id=project.id,
        status="queued",
        created_at=project.created_at,
    )


@router.get(
    "/jobs/{job_id}",
    response_model=MathRenderResponse,
    summary="Get status or download URL of a math video generation job",
)
def get_math_job_status(
    job_id: UUID,
    user_id: UUID = Depends(resolve_math_caller_user_id),  # noqa: B008
    content: ContentStore = Depends(get_content_store),  # noqa: B008
    hitl_store: SupabaseHitlStore = Depends(get_hitl_store),  # noqa: B008
    job_store: RedisRenderJobStore = Depends(get_job_store),  # noqa: B008
    pipeline_service: HitlPipelineService = Depends(get_pipeline_service),  # noqa: B008
) -> MathRenderResponse:
    """Lookup the status of a math rendering job by run_id or project_id."""
    # Try finding run first
    run = hitl_store.get_run(job_id)
    project_id = run.project_id if run else job_id

    project = content.get_project(project_id)
    if project is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Math generation job / project not found",
        )

    # 1. Check if the project already has a completed video
    if project.video_url:
        video_url = _resolve_video_url(project.id, project.video_url)
        return MathRenderResponse(
            job_id=job_id,
            project_id=project.id,
            status="completed",
            video_url=video_url,
            created_at=project.created_at,
            completed_at=project.updated_at,
        )

    # 2. Check active or finished render jobs
    render_jobs = job_store.list_for_project(project.id, active_only=False)
    for rjob in render_jobs:
        if rjob.status == "completed" and rjob.asset_url:
            video_url = _resolve_video_url(project.id, rjob.asset_url)
            return MathRenderResponse(
                job_id=job_id,
                project_id=project.id,
                status="completed",
                video_url=video_url,
                created_at=rjob.created_at,
                completed_at=rjob.completed_at,
            )
        if rjob.status in {"queued", "rendering"}:
            return MathRenderResponse(
                job_id=job_id,
                project_id=project.id,
                status="rendering",
                created_at=rjob.created_at,
            )
        if rjob.status == "failed":
            return MathRenderResponse(
                job_id=job_id,
                project_id=project.id,
                status="failed",
                error=rjob.logs or "Render failed",
                created_at=rjob.created_at,
                completed_at=rjob.completed_at,
            )

    # 3. Check AI runs status and expire any dead/timed-out tasks
    runs = [run] if run else hitl_store.list_runs(project.id)
    for r in runs:
        for s in hitl_store.list_steps(r.id):
            pipeline_service.expire_stale_generation(run=r, step=s)

    if runs and any(r.status == "failed" for r in runs):
        failed_run = next(r for r in runs if r.status == "failed")
        steps = hitl_store.list_steps(failed_run.id)
        failed_step = next((s for s in steps if s.status == "failed"), None)
        error_msg = failed_step.error if failed_step else "AI generation failed"
        return MathRenderResponse(
            job_id=job_id,
            project_id=project.id,
            status="failed",
            error=error_msg,
            created_at=project.created_at,
        )

    # 4. Check if all scenes are completed and need automatic project render
    scenes = content.get_project_scenes(project.id)
    if scenes and all(scene.generation_status == "completed" and scene.manim_code for scene in scenes):
        # Trigger full project render automatically
        quality_pref: RenderQuality = (project.config or {}).get("quality", "720p")
        render_meta = project_render_source(scenes)
        try:
            new_job_id = uuid4()
            rjob, created = job_store.get_or_create_active_job(
                job_id=new_job_id,
                project_id=project.id,
                scene_id=None,
                job_type="full_project",
                render_quality=quality_pref,
                docker_image_tag=None,
                metadata=render_meta,
            )
            if created:
                publish_project_event(
                    str(project.id),
                    "render.queued",
                    {"job_id": str(rjob.id), "job": rjob.model_dump(mode="json"), "scene_id": None},
                )
                AiQueue().dispatch_render(rjob.id)
                logger.info("Auto-dispatched project render for math project_id=%s", project.id)
            return MathRenderResponse(
                job_id=job_id,
                project_id=project.id,
                status="rendering",
                created_at=project.created_at,
            )
        except (AiQueueUnavailable, Exception) as exc:  # noqa: BLE001
            logger.warning("Auto project render dispatch failed for %s: %s", project.id, exc)

    return MathRenderResponse(
        job_id=job_id,
        project_id=project.id,
        status="generating",
        created_at=project.created_at,
    )