"""MiniMax H3 video client for the New API gateway."""

import json
from typing import Any, Callable, Dict, Optional
from urllib.parse import quote

import aiohttp

from ..utils.config import get_api_key_or_raise
from ..utils.http_error import async_request_with_retry
from ..utils.video_task import (
    POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
    PollDeadline,
    check_interrupt,
    download_video_to_file,
    extract_error_message,
    extract_progress,
    interruptible_sleep,
    run_with_interrupt,
)


PENDING_STATUSES = {
    "NOT_START",
    "SUBMITTED",
    "QUEUED",
    "IN_PROGRESS",
    "RUNNING",
    "UNKNOWN",
}
SUCCESS_STATUSES = {"SUCCESS", "COMPLETED", "SUCCEEDED"}
FAILURE_STATUSES = {"FAILURE", "FAILED", "CANCELLED", "CANCELED"}


def extract_public_task_id(payload: Dict[str, Any]) -> str:
    """Return the New API public task ID, preferring ``id`` as documented."""
    task_id = payload.get("id") or payload.get("task_id")
    if not task_id:
        raise RuntimeError("MiniMax H3 创建成功但未返回任务 ID。")
    return str(task_id)


def parse_task_snapshot(payload: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize New API's wrapper and MiniMax's official V2 task shape."""
    if not isinstance(payload, dict):
        raise RuntimeError("MiniMax H3 查询响应不是 JSON 对象。")

    raw_data = payload.get("data")
    data = raw_data if isinstance(raw_data, dict) else payload
    raw_task = payload.get("task")
    task = raw_task if isinstance(raw_task, dict) else {}
    task_content = task.get("content") if isinstance(task.get("content"), dict) else {}
    task_error = task.get("error") if isinstance(task.get("error"), dict) else {}
    data_error = data.get("error") if isinstance(data.get("error"), dict) else {}
    root_error = payload.get("error") if isinstance(payload.get("error"), dict) else {}
    metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {}
    status = str(
        data.get("status") or task.get("status") or payload.get("status") or ""
    ).strip().upper()

    error_message = str(
        task_error.get("message")
        or data_error.get("message")
        or root_error.get("message")
        or ""
    ).strip()
    error_code = str(
        task_error.get("code")
        or data_error.get("code")
        or root_error.get("code")
        or ""
    ).strip()
    if error_message and error_code:
        error_message = f"{error_message}（错误码 {error_code}）"
    elif error_code:
        error_message = f"错误码 {error_code}"

    return {
        "status": status,
        "progress": extract_progress(payload),
        "result_url": str(
            data.get("result_url")
            or task_content.get("url")
            or metadata.get("url")
            or data.get("url")
            or payload.get("result_url")
            or payload.get("url")
            or ""
        ).strip(),
        "fail_reason": str(
            data.get("fail_reason")
            or error_message
            or extract_error_message(payload, "视频生成失败")
        ).strip(),
    }


class MiniMaxH3Client:
    """Create, poll, and immediately download a MiniMax-H3 video task."""

    CREATE_ENDPOINT = "/v1/video/generations"
    STATUS_ENDPOINT = "/v1/videos/{task_id}"
    POLL_INTERVAL_SECONDS = 10.0
    POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS

    def __init__(self, base_url: str, api_key: Optional[str] = None):
        self.base_url = (base_url or "").rstrip("/")
        if not self.base_url:
            raise ValueError("MiniMax H3 New API Base URL 不能为空。")
        self.api_key = api_key or get_api_key_or_raise()

    def _headers(self) -> Dict[str, str]:
        # Model API requests use the application API token. New-Api-User is
        # intentionally not sent because it belongs to management API auth.
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    @staticmethod
    async def _read_json(response: aiohttp.ClientResponse, action: str) -> Dict[str, Any]:
        raw = await response.text()
        try:
            payload = json.loads(raw)
        except json.JSONDecodeError:
            raise RuntimeError(f"MiniMax H3 {action}返回了无效 JSON。") from None
        if not isinstance(payload, dict):
            raise RuntimeError(f"MiniMax H3 {action}响应不是 JSON 对象。")
        return payload

    async def submit_async(
        self,
        body: Dict[str, Any],
        session: aiohttp.ClientSession,
    ) -> str:
        url = f"{self.base_url}{self.CREATE_ENDPOINT}"
        check_interrupt()
        response = await run_with_interrupt(
            async_request_with_retry(
                session,
                "POST",
                url,
                json=body,
                headers=self._headers(),
                prefix="MiniMax H3 创建任务：",
            )
        )
        payload = await self._read_json(response, "创建任务")
        return extract_public_task_id(payload)

    async def poll_async(
        self,
        task_id: str,
        session: aiohttp.ClientSession,
        on_progress: Optional[Callable[[int], None]] = None,
    ) -> str:
        encoded_task_id = quote(task_id, safe="")
        url = f"{self.base_url}{self.STATUS_ENDPOINT.format(task_id=encoded_task_id)}"
        deadline = PollDeadline(
            seconds=self.POLL_DEADLINE_SECONDS,
            label=f"MiniMax H3（任务 {task_id}）",
        )

        while True:
            deadline.check()
            check_interrupt()
            response = await run_with_interrupt(
                async_request_with_retry(
                    session,
                    "GET",
                    url,
                    headers=self._headers(),
                    prefix="MiniMax H3 查询任务：",
                )
            )
            payload = await self._read_json(response, "查询任务")
            snapshot = parse_task_snapshot(payload)
            status = snapshot["status"]
            progress = snapshot["progress"]

            # A successful terminal state is authoritative even if an older
            # gateway omits data.progress or returns a stale percentage.
            if status in SUCCESS_STATUSES:
                progress = 100

            print(f"[MiniMax H3] 任务 {task_id}：{status or 'UNKNOWN'} {progress}%")
            if on_progress:
                on_progress(progress)

            if status in SUCCESS_STATUSES:
                result_url = snapshot["result_url"]
                if not result_url:
                    raise RuntimeError(
                        f"MiniMax H3 任务 {task_id} 已成功，但响应缺少 data.result_url。"
                    )
                return result_url

            if status in FAILURE_STATUSES:
                raise RuntimeError(
                    f"MiniMax H3 任务 {task_id} 生成失败：{snapshot['fail_reason']}"
                )

            if status not in PENDING_STATUSES:
                raise RuntimeError(
                    f"MiniMax H3 任务 {task_id} 返回不支持的状态 {status or '<空>'}。"
                )

            await interruptible_sleep(self.POLL_INTERVAL_SECONDS)

    async def generate_async(
        self,
        body: Dict[str, Any],
        save_path: str,
        on_stage: Optional[Callable[[str], None]] = None,
        on_progress: Optional[Callable[[int], None]] = None,
    ) -> tuple[str, str]:
        connector = aiohttp.TCPConnector(force_close=True)
        async with aiohttp.ClientSession(connector=connector) as session:
            if on_stage:
                on_stage("submitting")
            task_id = await self.submit_async(body, session)
            print(f"[MiniMax H3] 已提交公开任务 ID：{task_id}")
            if on_stage:
                on_stage(f"submitted:{task_id}")

            result_url = await self.poll_async(
                task_id,
                session,
                on_progress=on_progress,
            )

            if on_stage:
                on_stage("downloading")
            await download_video_to_file(
                session,
                result_url,
                save_path,
                label=f"MiniMax H3 {task_id}",
            )
            if on_stage:
                on_stage("done")
            return save_path, task_id