app/tasks.py

"""Celery Tasks & Async Worker Handlers for MathSolver Solve & Render Pipeline."""

from __future__ import annotations

import asyncio
import logging
import os
import uuid
from typing import Any, Dict, Optional

from app.celery_app import celery_app
from app.errors import format_error_for_user
from app.logutil import log_pipeline_failure, log_pipeline_success, log_step
from app.models.job_state import JobStatus, JobStage, JobStateMachine
from app.supabase_client import get_supabase
from app.websocket_manager import notify_status

logger = logging.getLogger(__name__)


async def async_solve_session_job(
    job_id: str,
    session_id: str,
    text: str,
    image_url: Optional[str] = None,
    user_id: Optional[str] = None,
    client_message_id: Optional[str] = None,
) -> Dict[str, Any]:
    """Execute the full geometry solve pipeline for a session job."""
    from app.routers.solve import get_orchestrator

    supabase = get_supabase()

    async def status_callback(status: str, stage: Optional[str] = None, progress: Optional[int] = None):
        norm_status = JobStateMachine.normalize_status(status)
        norm_stage = JobStateMachine.normalize_stage(stage or status)
        update_data = {"status": norm_status.value}

        if supabase:
            try:
                supabase.table("jobs").update(update_data).eq("id", job_id).execute()
            except Exception as e:
                logger.debug("Failed updating job status in DB: %s", e)

        await notify_status(job_id, {
            "status": norm_status.value,
            "stage": norm_stage.value if norm_stage else None,
            "progress": progress,
            "job_id": job_id,
        })

    try:
        # Initial status update
        await status_callback("processing", stage="ocr", progress=15)

        history = []
        if supabase and session_id:
            try:
                history_res = (
                    supabase.table("messages")
                    .select("*")
                    .eq("session_id", session_id)
                    .order("created_at", desc=False)
                    .execute()
                )
                history = history_res.data if history_res.data else []
            except Exception as e:
                logger.warning("Could not fetch message history: %s", e)

        result = await get_orchestrator().run(
            text,
            image_url,
            job_id=job_id,
            session_id=session_id,
            status_callback=lambda st: status_callback(st),
            history=history,
        )

        has_error = "error" in result and result.get("error")
        final_status = JobStatus.FAILED if has_error else JobStatus.COMPLETED

        if supabase:
            supabase.table("jobs").update({
                "status": final_status.value,
                "result": result,
            }).eq("id", job_id).execute()

            # Idempotency check: Ensure assistant message for this job is not inserted twice
            existing_msg = (
                supabase.table("messages")
                .select("id")
                .eq("session_id", session_id)
                .filter("metadata->>job_id", "eq", job_id)
                .execute()
            )

            if not existing_msg.data or len(existing_msg.data) == 0:
                supabase.table("messages").insert({
                    "session_id": session_id,
                    "role": "assistant",
                    "type": "error" if has_error else "analysis",
                    "content": (
                        result.get("error", "Đã có lỗi xảy ra.")
                        if has_error
                        else result.get("semantic_analysis", "Giải bài toán hoàn tất.")
                    ),
                    "metadata": {
                        "job_id": job_id,
                        "client_message_id": client_message_id,
                        "coordinates": result.get("coordinates"),
                        "geometry_dsl": result.get("geometry_dsl"),
                        "polygon_order": result.get("polygon_order", []),
                        "drawing_phases": result.get("drawing_phases", []),
                        "circles": result.get("circles", []),
                        "solids": result.get("solids", []),
                        "faces": result.get("faces", []),
                        "lines": result.get("lines", []),
                        "rays": result.get("rays", []),
                        "visualization_graph": result.get("visualization_graph"),
                        "auxiliary": result.get("auxiliary", []),
                        "solution": result.get("solution"),
                        "is_3d": result.get("is_3d", False),
                    },
                }).execute()

        await notify_status(job_id, {
            "status": final_status.value,
            "stage": None,
            "progress": 100 if final_status == JobStatus.COMPLETED else 0,
            "job_id": job_id,
            "result": result,
        })
        log_pipeline_success("job_complete", job_id=job_id, session_id=session_id)
        return result

    except Exception as e:
        logger.exception("Error in async_solve_session_job for job %s: %s", job_id, e)
        error_msg = format_error_for_user(e)
        if supabase:
            try:
                supabase.table("jobs").update({
                    "status": JobStatus.FAILED.value,
                    "result": {"error": str(e)},
                }).eq("id", job_id).execute()

                supabase.table("messages").insert({
                    "session_id": session_id,
                    "role": "assistant",
                    "type": "error",
                    "content": error_msg,
                    "metadata": {"job_id": job_id, "client_message_id": client_message_id},
                }).execute()
            except Exception as dbe:
                logger.error("DB error recording failure for job %s: %s", job_id, dbe)

        await notify_status(job_id, {
            "status": JobStatus.FAILED.value,
            "job_id": job_id,
            "error": error_msg,
            "progress": 0,
        })
        log_pipeline_failure("job_failed", job_id=job_id, error=str(e))
        return {"status": "error", "error": error_msg}


async def async_render_video_job(job_id: str, session_id: str, geometry_data: Dict[str, Any]) -> Dict[str, Any]:
    """Execute Manim video rendering job for a session."""
    from manim_client import ManimClient, build_visualization_spec
    from manim_client.schemas import ErrorCode

    await notify_status(job_id, {
        "status": JobStatus.QUEUED.value,
        "stage": JobStage.RENDERING.value,
        "job_id": job_id,
        "progress": 10,
    })
    supabase = get_supabase()

    try:
        manim_url = os.getenv("MANIM_SERVICE_URL", "https://cuong2004-manim-agent.hf.space")
        manim_token = os.getenv("MANIM_INTERNAL_TOKEN")
        client = ManimClient(base_url=manim_url, internal_token=manim_token)

        vis_spec = build_visualization_spec(geometry_data)
        resp = await client.submit_render_job(vis_spec)

        if resp.status == "failed":
            err_code = resp.get_error_code() or ErrorCode.MANIM_REQUEST_FAILED
            err_msg = resp.get_error_message() or "Không thể gửi yêu cầu tạo video tới máy chủ Manim."
            if supabase:
                supabase.table("jobs").update({
                    "status": JobStatus.FAILED.value,
                    "result": {"error": {"code": err_code, "message": err_msg}},
                }).eq("id", job_id).execute()
                if session_id:
                    supabase.table("messages").insert({
                        "session_id": session_id,
                        "role": "assistant",
                        "type": "error",
                        "content": f"Không thể tạo video: {err_msg}",
                        "metadata": {"job_id": job_id, "error_code": err_code},
                    }).execute()
            await notify_status(job_id, {
                "status": JobStatus.FAILED.value,
                "job_id": job_id,
                "error": err_msg,
                "error_code": err_code,
            })
            return {"status": "error", "error": err_msg}

        manim_job_id = resp.job_id
        if supabase:
            supabase.table("jobs").update({
                "status": JobStatus.PROCESSING.value,
                "result": {"manim_job_id": str(manim_job_id)},
            }).eq("id", job_id).execute()

        await notify_status(job_id, {
            "status": JobStatus.PROCESSING.value,
            "stage": JobStage.RENDERING.value,
            "job_id": job_id,
            "progress": 40,
            "manim_job_id": str(manim_job_id),
        })

        poll_timeout = float(os.getenv("MANIM_POLL_TIMEOUT", "600.0"))
        status_resp = await client.wait_for_completion(manim_job_id, poll_interval=3.0, max_wait=poll_timeout)
        video_url = status_resp.video_url

        if status_resp.status == "failed" or not video_url:
            err_code = status_resp.get_error_code() or ErrorCode.MANIM_RENDER_FAILED
            err_msg = status_resp.get_error_message() or "Tiến trình dựng video Manim thất bại."
            if supabase:
                supabase.table("jobs").update({
                    "status": JobStatus.FAILED.value,
                    "result": {"error": {"code": err_code, "message": err_msg}},
                }).eq("id", job_id).execute()
                if session_id:
                    supabase.table("messages").insert({
                        "session_id": session_id,
                        "role": "assistant",
                        "type": "error",
                        "content": f"Không thể tạo video: {err_msg}",
                        "metadata": {"job_id": job_id, "error_code": err_code},
                    }).execute()
            await notify_status(job_id, {
                "status": JobStatus.FAILED.value,
                "job_id": job_id,
                "error": err_msg,
                "error_code": err_code,
            })
            return {"status": "error", "error": err_msg}

        final_result = geometry_data.copy()
        final_result["video_url"] = video_url
        final_result["manim_job_id"] = str(manim_job_id)

        if supabase:
            supabase.table("jobs").update({
                "status": JobStatus.COMPLETED.value,
                "result": final_result,
            }).eq("id", job_id).execute()

            # Versioned asset recording
            try:
                asset_version = 1
                v_res = supabase.table("session_assets").select("version").eq("session_id", session_id).eq("asset_type", "video").order("version", desc=True).limit(1).execute()
                if v_res.data and len(v_res.data) > 0:
                    asset_version = v_res.data[0]["version"] + 1

                supabase.table("session_assets").insert({
                    "session_id": session_id,
                    "job_id": job_id,
                    "asset_type": "video",
                    "storage_path": video_url,
                    "public_url": video_url,
                    "version": asset_version,
                }).execute()
            except Exception as e:
                logger.warning("Could not record session_asset video row: %s", e)

            if session_id:
                supabase.table("messages").insert({
                    "session_id": session_id,
                    "role": "assistant",
                    "type": "analysis",
                    "content": geometry_data.get("semantic_analysis", "🎬 Video minh họa hình học đã hoàn tất."),
                    "metadata": {
                        "job_id": job_id,
                        "video_url": video_url,
                        "coordinates": geometry_data.get("coordinates"),
                        "geometry_dsl": geometry_data.get("geometry_dsl"),
                        "polygon_order": geometry_data.get("polygon_order", []),
                        "drawing_phases": geometry_data.get("drawing_phases", []),
                        "circles": geometry_data.get("circles", []),
                        "solids": geometry_data.get("solids", []),
                        "faces": geometry_data.get("faces", []),
                        "lines": geometry_data.get("lines", []),
                        "rays": geometry_data.get("rays", []),
                        "visualization_graph": geometry_data.get("visualization_graph"),
                        "auxiliary": geometry_data.get("auxiliary", []),
                        "is_3d": geometry_data.get("is_3d", False),
                    },
                }).execute()

        await notify_status(job_id, {
            "status": JobStatus.COMPLETED.value,
            "job_id": job_id,
            "result": final_result,
            "video_url": video_url,
            "progress": 100,
        })
        return final_result

    except Exception as e:
        logger.exception("Error rendering video for job %s: %s", job_id, e)
        safe_msg = format_error_for_user(e)
        if supabase:
            try:
                supabase.table("jobs").update({
                    "status": JobStatus.FAILED.value,
                    "result": {"error": {"message": safe_msg}},
                }).eq("id", job_id).execute()
                if session_id:
                    supabase.table("messages").insert({
                        "session_id": session_id,
                        "role": "assistant",
                        "type": "error",
                        "content": f"Lỗi render video: {safe_msg}",
                        "metadata": {"job_id": job_id},
                    }).execute()
            except Exception as dbe:
                logger.error("DB error recording render failure: %s", dbe)
        await notify_status(job_id, {"status": JobStatus.FAILED.value, "job_id": job_id, "error": safe_msg})
        return {"status": "error", "error": safe_msg}


@celery_app.task(name="tasks.solve_session_job", bind=True, acks_late=True, max_retries=1)
def solve_session_job_task(
    self,
    job_id: str,
    session_id: str,
    text: str,
    image_url: Optional[str] = None,
    user_id: Optional[str] = None,
    client_message_id: Optional[str] = None,
):
    """Celery task entry point for solve pipeline."""
    return asyncio.run(
        async_solve_session_job(
            job_id=job_id,
            session_id=session_id,
            text=text,
            image_url=image_url,
            user_id=user_id,
            client_message_id=client_message_id,
        )
    )


@celery_app.task(name="tasks.render_video_job", bind=True, acks_late=True, max_retries=1)
def render_video_job_task(self, job_id: str, session_id: str, geometry_data: Dict[str, Any]):
    """Celery task entry point for video render pipeline."""
    return asyncio.run(async_render_video_job(job_id=job_id, session_id=session_id, geometry_data=geometry_data))


def recover_stale_jobs(timeout_seconds: int = 900) -> int:
    """Detect and mark jobs stuck in 'processing' longer than timeout as failed."""
    supabase = get_supabase()
    if not supabase:
        return 0
    try:
        # Note: in production, run via cron or worker startup
        from datetime import datetime, timezone, timedelta
        cutoff = (datetime.now(timezone.utc) - timedelta(seconds=timeout_seconds)).isoformat()
        res = (
            supabase.table("jobs")
            .update({
                "status": JobStatus.FAILED.value,
                "result": {"error": "Worker timeout or crash detected. Job marked failed by recovery agent."},
            })
            .eq("status", JobStatus.PROCESSING.value)
            .lt("created_at", cutoff)
            .execute()
        )
        count = len(res.data) if res.data else 0
        if count > 0:
            logger.warning("Recovered %d stale jobs", count)
        return count
    except Exception as e:
        logger.error("Error running recover_stale_jobs: %s", e)
        return 0