manim_client/client.py

"""Async Client for interacting with the external Manim Video Generation Module."""
from __future__ import annotations

import asyncio
import logging
import os
from typing import Any, Dict, Optional, Union
from uuid import UUID
import httpx
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)

from manim_client.schemas import (
    ErrorCode,
    MathRenderRequest,
    MathRenderResponse,
    StructuredError,
    VisualizationSpec,
)

DEFAULT_MANIM_SERVICE_URL = os.getenv("MANIM_SERVICE_URL", "http://127.0.0.1:8001")
DEFAULT_INTERNAL_TOKEN = os.getenv(
    "MANIM_INTERNAL_TOKEN",
    "4f8a3c2e1d0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b1a0f9e8d7c6b5a4",
)


class ManimClient:
    """
    Client connecting MathSolver to the external Manim Video Generation Agent.

    API Boundary:
    - POST /v1/math/generate: Submit VisualizationSpec
    - GET  /v1/math/jobs/{job_id}: Poll video rendering status and get video_url
    """

    def __init__(
        self,
        base_url: Optional[str] = None,
        internal_token: Optional[str] = None,
        timeout: float = 30.0,
    ):
        self.base_url = (base_url or DEFAULT_MANIM_SERVICE_URL).rstrip("/")
        self.internal_token = internal_token or DEFAULT_INTERNAL_TOKEN
        self.timeout = timeout

    def _headers(self) -> Dict[str, str]:
        return {
            "Content-Type": "application/json",
            "X-Internal-Token": self.internal_token,
        }

    async def check_health(self) -> bool:
        """Checks if the Manim video generation service is reachable."""
        try:
            async with httpx.AsyncClient(timeout=3.0) as client:
                resp = await client.get(f"{self.base_url}/health")
                return resp.status_code == 200
        except Exception:
            return False

    async def submit_render_job(
        self,
        spec: VisualizationSpec,
        callback_url: Optional[str] = None,
    ) -> MathRenderResponse:
        """
        Submits a VisualizationSpec to the Manim service to queue video generation.
        Endpoint: POST /v1/math/generate
        """
        url = f"{self.base_url}/v1/math/generate"
        spec_payload = spec.to_manim_dict() if hasattr(spec, "to_manim_dict") else spec.model_dump(mode="json")
        payload = {"spec": spec_payload, "callback_url": callback_url}

        logger.info(f"==[ManimClient] Submitting render job to {url}==")
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    url,
                    json=payload,
                    headers=self._headers(),
                )

                if response.status_code in (200, 201, 202):
                    data = response.json()
                    job_id = data.get("job_id")
                    if not job_id:
                        return MathRenderResponse(
                            job_id="error",
                            status="failed",
                            error=StructuredError(
                                code=ErrorCode.MANIM_REQUEST_FAILED,
                                message="Manim service did not return a valid job_id.",
                            ),
                        )
                    logger.info(
                        f"[ManimClient] Render job queued successfully: job_id={job_id}, status={data.get('status')}"
                    )
                    return MathRenderResponse.model_validate(data)
                else:
                    error_msg = f"HTTP {response.status_code}: {response.text[:200]}"
                    logger.warning(f"[ManimClient] Error response from Manim service: {error_msg}")
                    return MathRenderResponse(
                        job_id="error",
                        status="failed",
                        error=StructuredError(
                            code=ErrorCode.MANIM_REQUEST_FAILED,
                            message=f"Dịch vụ tạo video trả về mã lỗi HTTP {response.status_code}.",
                        ),
                    )
        except (httpx.ConnectError, httpx.ConnectTimeout) as e:
            logger.warning(f"[ManimClient] Connection to Manim service at {self.base_url} failed: {e}")
            return MathRenderResponse(
                job_id="offline",
                status="failed",
                error=StructuredError(
                    code=ErrorCode.MANIM_UNAVAILABLE,
                    message="Không thể kết nối đến dịch vụ tạo video Manim (máy chủ ngoại vi không khả dụng).",
                ),
            )
        except httpx.TimeoutException as e:
            logger.warning(f"[ManimClient] Request to Manim service timed out: {e}")
            return MathRenderResponse(
                job_id="timeout",
                status="failed",
                error=StructuredError(
                    code=ErrorCode.MANIM_TIMEOUT,
                    message="Yêu cầu gửi sang dịch vụ tạo video đã hết thời gian chờ (request timeout).",
                ),
            )
        except Exception as e:
            logger.exception(f"[ManimClient] Unexpected error submitting render job: {e}")
            return MathRenderResponse(
                job_id="error",
                status="failed",
                error=StructuredError(
                    code=ErrorCode.INTERNAL_ERROR,
                    message="Đã xảy ra lỗi không xác định khi yêu cầu tạo video.",
                ),
            )

    async def get_job_status(self, job_id: Union[UUID, str]) -> MathRenderResponse:
        """
        Polls the status of a video generation job.
        Endpoint: GET /v1/math/jobs/{job_id}
        """
        url = f"{self.base_url}/v1/math/jobs/{job_id}"
        logger.debug(f"[ManimClient] Fetching status for job {job_id}")

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.get(url, headers=self._headers())
                if response.status_code == 200:
                    data = response.json()
                    resp = MathRenderResponse.model_validate(data)
                    if resp.status == "failed" and not isinstance(resp.error, StructuredError):
                        raw_err = resp.get_error_message() or "Animation rendering failed."
                        resp.error = StructuredError(
                            code=ErrorCode.MANIM_RENDER_FAILED,
                            message=raw_err,
                        )
                    return resp
                elif response.status_code == 404:
                    return MathRenderResponse(
                        job_id=job_id,
                        status="failed",
                        error=StructuredError(
                            code=ErrorCode.JOB_NOT_FOUND,
                            message=f"Không tìm thấy tiến trình render video với ID '{job_id}'.",
                        ),
                    )
                else:
                    return MathRenderResponse(
                        job_id=job_id,
                        status="failed",
                        error=StructuredError(
                            code=ErrorCode.MANIM_REQUEST_FAILED,
                            message=f"Lỗi kiểm tra tiến trình: HTTP {response.status_code}.",
                        ),
                    )
        except (httpx.ConnectError, httpx.ConnectTimeout) as e:
            return MathRenderResponse(
                job_id=job_id,
                status="failed",
                error=StructuredError(
                    code=ErrorCode.MANIM_UNAVAILABLE,
                    message="Không thể kết nối đến máy chủ render video để kiểm tra trạng thái.",
                ),
            )
        except Exception as e:
            return MathRenderResponse(
                job_id=job_id,
                status="failed",
                error=StructuredError(
                    code=ErrorCode.INTERNAL_ERROR,
                    message=f"Lỗi kiểm tra trạng thái: {str(e)}",
                ),
            )

    async def poll_job_completion(
        self,
        job_id: Union[UUID, str],
        timeout: float = 300.0,
        poll_interval: float = 3.0,
    ) -> MathRenderResponse:
        """
        Asynchronously polls until the job reaches a terminal status ('completed' or 'failed') or times out.
        """
        start_time = asyncio.get_event_loop().time()
        while True:
            resp = await self.get_job_status(job_id)
            if resp.is_terminal():
                return resp

            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed >= timeout:
                logger.warning(f"[ManimClient] Polling for job {job_id} timed out after {timeout:.1f}s")
                return MathRenderResponse(
                    job_id=job_id,
                    status="failed",
                    error=StructuredError(
                        code=ErrorCode.MANIM_TIMEOUT,
                        message=f"Tiến trình dựng video đã vượt quá thời gian tối đa ({int(timeout)} giây).",
                    ),
                )

            await asyncio.sleep(poll_interval)

    async def wait_for_completion(
        self,
        job_id: Union[UUID, str],
        poll_interval: float = 3.0,
        max_wait: float = 300.0,
    ) -> MathRenderResponse:
        """Alias for poll_job_completion for API compatibility."""
        return await self.poll_job_completion(job_id=job_id, timeout=max_wait, poll_interval=poll_interval)