"""Client for the complete O1Key Grok Imagine Video API."""

import asyncio
import json
import os
import re
import time
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import quote

import aiohttp

from .base_client import BaseAPIClient
from ..utils.config import get_api_base_url, get_api_key_or_raise
from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message
from ..utils.video_task import (
    POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
    check_interrupt,
    download_video_to_file,
    extract_error_message,
    interruptible_sleep,
    run_with_interrupt,
)


class GrokVideoClient(BaseAPIClient):
    """Submit, poll, and download Grok video generation, edit, or extension tasks."""

    ENDPOINTS = {
        "generate": "/grok/v1/videos/generations",
        "edit": "/grok/v1/videos/edits",
        "extend": "/grok/v1/videos/extensions",
    }
    STATUS_ENDPOINT = "/grok/v1/videos/{request_id}"

    BASE_MODEL = "grok-imagine-video"
    LATEST_MODEL = "grok-imagine-video-1.5"
    DEFAULT_MODEL = LATEST_MODEL
    # Kept as a compatibility alias for callers that imported the old constant.
    TEXT_TO_VIDEO_MODEL = BASE_MODEL
    IMAGE_TO_VIDEO_MODELS = (LATEST_MODEL,)
    MODEL_OPTIONS = (BASE_MODEL, LATEST_MODEL)
    ASPECT_RATIO_OPTIONS = ("16:9", "9:16", "1:1", "4:3", "3:4", "3:2", "2:3")
    RESOLUTION_OPTIONS = ("480p", "720p", "1080p")
    SUCCESS_STATUSES = {"done"}
    FAILURE_STATUSES = {"failed", "expired"}
    POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS

    def __init__(self, base_url: Optional[str] = None):
        super().__init__(
            base_url=(base_url or get_api_base_url()).rstrip("/"),
            api_key=get_api_key_or_raise("O1KEY_API_KEY"),
        )

    def get_endpoint(self, operation: str = "generate", **kwargs) -> str:
        try:
            return self.ENDPOINTS[operation]
        except KeyError:
            raise ValueError(f"不支持的 Grok 操作：{operation}。") from None

    def build_request_body(self, **kwargs) -> Dict[str, Any]:
        return self.build_video_body(**kwargs)

    def parse_response(self, response: Dict[str, Any]) -> Any:
        return response

    @staticmethod
    def _locator(
        value: Optional[Dict[str, str]],
        label: str,
        allowed_keys: tuple[str, ...],
    ) -> Dict[str, str]:
        if not isinstance(value, dict):
            raise ValueError(f"{label}必须提供媒体定位对象。")
        known_keys = ("url", "image_url", "file_id", "voice_id")
        provided_keys = {
            key
            for key in known_keys
            if value.get(key) is not None and str(value[key]).strip()
        }
        locator = {
            key: str(value[key]).strip()
            for key in allowed_keys
            if value.get(key) is not None and str(value[key]).strip()
        }
        if len(locator) != 1 or provided_keys != set(locator):
            supported = "、".join(allowed_keys)
            raise ValueError(f"{label}必须且只能提供 {supported} 中的一项。")
        return locator

    @classmethod
    def _image_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
        return cls._locator(value, label, ("url", "image_url"))

    @classmethod
    def _audio_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
        return cls._locator(value, label, ("url", "voice_id"))

    @classmethod
    def _video_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
        return cls._locator(value, label, ("url", "file_id"))

    @classmethod
    def _validate_common_generation(
        cls, model: str, duration: int, aspect_ratio: str, resolution: str
    ) -> int:
        if model not in cls.MODEL_OPTIONS:
            raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
        try:
            duration = int(duration)
        except (TypeError, ValueError):
            raise ValueError("时长必须是整数。") from None
        if not 1 <= duration <= 15:
            raise ValueError("生成时长仅支持 1 到 15 秒。")
        if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
            raise ValueError(f"宽高比仅支持：{', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
        if resolution not in cls.RESOLUTION_OPTIONS:
            raise ValueError(f"分辨率仅支持：{', '.join(cls.RESOLUTION_OPTIONS)}。")
        return duration

    @classmethod
    def build_video_body(
        cls,
        *,
        operation: str,
        prompt: str,
        model: str,
        duration: Optional[int] = None,
        aspect_ratio: str = "16:9",
        resolution: str = "480p",
        image: Optional[Dict[str, str]] = None,
        reference_images: Optional[List[Dict[str, str]]] = None,
        reference_audios: Optional[List[Dict[str, str]]] = None,
        video: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        if operation not in cls.ENDPOINTS:
            raise ValueError(f"不支持的 Grok 操作：{operation}。")
        prompt = (prompt or "").strip()
        if reference_images is not None and not isinstance(reference_images, (list, tuple)):
            raise ValueError("reference_images 必须是数组。")
        if reference_audios is not None and not isinstance(reference_audios, (list, tuple)):
            raise ValueError("reference_audios 必须是数组。")
        references = list(reference_images or [])
        audios = list(reference_audios or [])

        if operation == "generate":
            duration = cls._validate_common_generation(model, duration, aspect_ratio, resolution)
            normal_image = cls._image_locator(image, "图生视频参考图") if image else None
            normal_references = [cls._image_locator(item, "参考图") for item in references]
            normal_audios = [cls._audio_locator(item, "参考音频") for item in audios]
            if normal_image and normal_references:
                raise ValueError("image 和 reference_images 不能同时使用。")
            if len(normal_references) > 7:
                raise ValueError("参考生视频最多支持 7 张参考图。")
            if len(normal_audios) > 3:
                raise ValueError("参考生视频最多支持 3 个参考音频。")

            has_reference_assets = bool(normal_references or normal_audios)
            if has_reference_assets:
                if not prompt:
                    raise ValueError("参考图/音频生视频必须填写提示词。")
                if resolution == "1080p":
                    raise ValueError("参考图/音频生视频不支持 1080p。")
            elif not normal_image and not prompt:
                raise ValueError("文生视频必须填写提示词。")

            if resolution == "1080p" and model != cls.LATEST_MODEL:
                raise ValueError("1080p 仅支持 grok-imagine-video-1.5 的文生或图生视频。")

            body: Dict[str, Any] = {
                "model": model,
                "duration": duration,
                "aspect_ratio": aspect_ratio,
                "resolution": resolution,
            }
            if prompt:
                body["prompt"] = prompt
            if normal_image:
                body["image"] = normal_image
            if normal_references:
                body["reference_images"] = normal_references
            if normal_audios:
                body["reference_audios"] = normal_audios
            return body

        if model not in cls.MODEL_OPTIONS:
            raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
        if not prompt:
            raise ValueError(f"{operation} 必须填写提示词。")
        normal_video = cls._video_locator(video, "输入视频")
        if operation == "edit":
            return {"model": model, "prompt": prompt, "video": normal_video}

        try:
            duration = int(duration)
        except (TypeError, ValueError):
            raise ValueError("续写时长必须是整数。") from None
        if not 2 <= duration <= 10:
            raise ValueError("视频续写时长仅支持 2 到 10 秒。")
        return {
            "model": model,
            "prompt": prompt,
            "video": normal_video,
            "duration": duration,
        }

    @staticmethod
    def _extract_request_id(payload: Dict[str, Any]) -> Optional[str]:
        for source in (payload, payload.get("data")):
            if isinstance(source, dict) and source.get("request_id"):
                return str(source["request_id"])
        return None

    @staticmethod
    def _safe_filename(request_id: str) -> str:
        return re.sub(r"[^A-Za-z0-9_.-]+", "_", request_id).strip("._") or "grok_video"

    @staticmethod
    def _safe_error_message(value: object) -> str:
        message = str(value or "").strip()
        message = re.sub(r"data:[^\s,;]+;base64,[A-Za-z0-9+/=_-]+", "<base64 omitted>", message)
        message = re.sub(r"https?://[^\s\"'<>]+", "<temporary URL omitted>", message)
        return message[:500]

    async def _request_json(
        self,
        method: str,
        endpoint: str,
        session: aiohttp.ClientSession,
        *,
        json_body: Optional[Dict[str, Any]] = None,
        timeout_seconds: int = 120,
        request_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        url = f"{self.base_url}{endpoint}"
        timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
        last_status, last_text = 0, ""
        for attempt in range(4):
            check_interrupt()
            response = None
            try:
                response = await run_with_interrupt(
                    session.request(
                        method, url, json=json_body,
                        headers=self.get_headers(use_bearer_token=True), timeout=timeout,
                    )
                )
                text = await run_with_interrupt(response.text())
                last_status, last_text = response.status, text
                if 200 <= response.status < 300:
                    try:
                        return json.loads(text) if text.strip() else {}
                    except json.JSONDecodeError:
                        raise RuntimeError("Grok Video 响应不是有效 JSON。") from None
                if response.status not in RETRYABLE_STATUS_CODES or attempt == 3:
                    break
                delay = min(2 ** attempt, 8)
                print(f"Grok Video：HTTP {response.status}，{delay}s 后重试…")
                await interruptible_sleep(delay)
            except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
                if attempt == 3:
                    raise RuntimeError(
                        f"Grok Video 网络错误：{type(exc).__name__}"
                    ) from None
                delay = min(2 ** attempt, 8)
                print(f"Grok Video：网络错误，{delay}s 后重试…")
                await interruptible_sleep(delay)
            finally:
                if response is not None:
                    response.release()

        message = self._safe_error_message(
            get_friendly_message(last_status, last_text) or "请求失败"
        )
        detail = f"Grok Video 请求失败：HTTP {last_status}，{message}"
        if request_id:
            detail += f"（request_id: {request_id}）"
        raise RuntimeError(detail)

    async def _poll(
        self, request_id: str, session: aiohttp.ClientSession, *, poll_interval: int,
        timeout: int, progress_callback: Optional[Callable[[int, str, float], None]],
    ) -> Dict[str, Any]:
        endpoint = self.STATUS_ENDPOINT.format(request_id=quote(request_id, safe=""))
        started_at = time.monotonic()
        while True:
            await interruptible_sleep(poll_interval)
            response = await self._request_json(
                "GET", endpoint, session, timeout_seconds=60, request_id=request_id
            )
            status = str(response.get("status", "")).strip().lower()
            try:
                progress = max(0, min(100, int(float(response.get("progress") or 0))))
            except (TypeError, ValueError):
                progress = 0
            elapsed = time.monotonic() - started_at
            if progress_callback:
                progress_callback(progress, status, elapsed)
            if status in self.SUCCESS_STATUSES:
                return response
            if status in self.FAILURE_STATUSES:
                message = self._safe_error_message(
                    extract_error_message(response, default="未知错误")
                )
                raise RuntimeError(
                    f"Grok Video 任务{status}（request_id: {request_id}）："
                    f"{message}"
                )
            if elapsed >= timeout:
                raise TimeoutError(
                    f"Grok Video 轮询超时（request_id: {request_id}，状态：{status or 'unknown'}）。"
                )

    def run_video_sync(
        self, *, operation: str, prompt: str, model: str, duration: Optional[int] = None,
        aspect_ratio: str = "16:9", resolution: str = "480p",
        image: Optional[Dict[str, str]] = None,
        reference_images: Optional[List[Dict[str, str]]] = None,
        reference_audios: Optional[List[Dict[str, str]]] = None,
        video: Optional[Dict[str, str]] = None, output_dir: Optional[str] = None,
        poll_interval: int = 5, timeout: int = VIDEO_POLL_DEADLINE_SECONDS,
        progress_callback: Optional[Callable[[int, str, float], None]] = None,
    ) -> Dict[str, Any]:
        async def run_request() -> Dict[str, Any]:
            async with self._make_session() as session:
                endpoint = self.get_endpoint(operation)
                body = self.build_video_body(
                    operation=operation, prompt=prompt, model=model, duration=duration,
                    aspect_ratio=aspect_ratio, resolution=resolution, image=image,
                    reference_images=reference_images, reference_audios=reference_audios,
                    video=video,
                )
                print(f"Grok Video：正在提交{operation}任务…")
                created = await self._request_json(
                    "POST", endpoint, session, json_body=body, timeout_seconds=180
                )
                request_id = self._extract_request_id(created)
                if not request_id:
                    raise RuntimeError("Grok Video 创建响应中没有 request_id。")
                print(f"Grok Video：任务已提交，request_id：{request_id}")
                completed = await self._poll(
                    request_id, session, poll_interval=max(1, int(poll_interval)),
                    timeout=timeout, progress_callback=progress_callback,
                )
                video_data = completed.get("video")
                video_url = video_data.get("url") if isinstance(video_data, dict) else None
                if not video_url:
                    raise RuntimeError(f"Grok Video 完成响应中没有 video.url（request_id: {request_id}）。")
                directory = output_dir or os.getcwd()
                os.makedirs(directory, exist_ok=True)
                save_path = os.path.join(directory, f"{self._safe_filename(request_id)}.mp4")
                print("Grok Video：视频生成完成，正在下载…")
                video_path = await download_video_to_file(session, video_url, save_path, label="Grok Video")
                return {
                    "request_id": request_id,
                    "video_path": video_path,
                    "duration": video_data.get("duration"),
                    "raw_json": {"create": created, "status": completed},
                }

        return self.run_async_in_thread(run_request())