"""O1Key Omni Flash JSON video tasks."""

from __future__ import annotations

import asyncio
import json
import re
import time
from typing import Any, Callable

import aiohttp

from ..utils.config import get_api_key_or_raise, get_base_url_by_route
from ..utils.video_task import (
    InterruptProcessingException, check_interrupt, download_video_to_file,
    extract_error_message, extract_progress, extract_video_url,
    interruptible_sleep, is_failure_status, is_success_status,
)


MODELS = {"omni_flash_8s", "omni_flash_10s", "omni_flash_abra_edit"}
RESOLUTIONS = {"720p", "1080p"}
RATIOS = {"16:9", "9:16"}
POLL_SECONDS = 7
POLL_DEADLINE_SECONDS = 2000
ERROR_HINTS = {
    "invalid_request": "请求参数有误，请检查模型、分辨率、宽高比和素材",
    "model_not_available": "当前模型不可用，请重新选择模型",
    "image_url_required_for_i2v": "参考图地址缺失或无效，请连接图片后重试",
    "invalid_api_key": "O1Key 令牌无效或已停用，请在令牌管理中更新",
    "insufficient_balance": "O1Key 余额不足",
    "task_not_found": "视频任务不存在或已失效",
    "rate_limit_exceeded": "请求过于频繁，请稍后重试",
}
HTTP_HINTS = {
    400: "请求参数错误", 401: "令牌验证失败", 402: "余额不足",
    404: "任务不存在", 429: "请求过于频繁",
}
SENSITIVE_RESPONSE_KEYS = {
    "authorization", "api_key", "apikey", "access_token", "refresh_token",
    "token", "secret", "password", "b64_json", "base64", "image_base64",
    "video_base64",
}
MAX_LOG_BODY = 16000


def build_video_body(
    *, model: str, prompt: str, resolution: str, aspect_ratio: str,
    mode: str, references: list[str] | None = None, source_video_url: str = "",
) -> dict[str, Any]:
    """Validate all scalar inputs before a paid request."""
    if model not in MODELS:
        raise ValueError("Omni Flash 模型无效")
    prompt = str(prompt or "").strip()
    if not prompt:
        raise ValueError("提示词不能为空")
    if len(prompt) > 20000:
        raise ValueError("提示词过长")
    if resolution not in RESOLUTIONS or aspect_ratio not in RATIOS:
        raise ValueError("分辨率或宽高比无效")
    references = list(references or [])
    if any(not isinstance(value, str) or len(value) > 4096 or not value.startswith(("https://", "http://")) for value in references):
        raise ValueError("参考图必须是 HTTP(S) 直链")
    body: dict[str, Any] = {
        "model": model, "prompt": prompt,
        "resolution": resolution, "aspect_ratio": aspect_ratio,
    }
    if mode == "edit":
        if model != "omni_flash_abra_edit" or len(source_video_url) > 4096 or not source_video_url.startswith(("https://", "http://")):
            raise ValueError("视频编辑需要编辑模型和源视频直链")
        if len(references) > 5:
            raise ValueError("视频编辑最多支持 5 张参考图")
        body["source_video_url"] = source_video_url
    elif mode in {"text", "reference", "first_last_frame"}:
        if model == "omni_flash_abra_edit" or source_video_url:
            raise ValueError("生成模式不能使用编辑模型或源视频")
        if mode == "text" and references:
            raise ValueError("文生视频不能提供参考图")
        if mode == "reference" and not references:
            raise ValueError("参考图模式至少需要 1 张图片")
        if mode == "first_last_frame":
            if not 1 <= len(references) <= 2:
                raise ValueError("首尾帧模式需要首帧图片，尾帧图片可选")
            # The provider's frame-pair flag is for a transition between two
            # frames. A lone first frame uses the documented single-image i2v
            # request, avoiding a pair request with a missing end frame.
            if len(references) == 2:
                body["first_last_frame"] = True
    else:
        raise ValueError("Omni Flash 生成模式无效")
    if references:
        body["input_reference"] = references[0] if len(references) == 1 else references
    return body


def _submission_payload(body: dict[str, Any]) -> dict[str, Any]:
    """Use a scalar JSON reference, or repeat the field in multipart for several."""
    references = body.get("input_reference")
    if not isinstance(references, list):
        return {"json": body}
    form = aiohttp.FormData()
    for name, value in body.items():
        values = value if name == "input_reference" else [value]
        for item in values:
            text = "true" if item is True else "false" if item is False else str(item)
            form.add_field(name, text, content_type="text/plain")
    return {"data": form}


def _redact_log_string(text: str) -> str:
    text = re.sub(r"https?://[^\s\"'<>]+", "<URL 已隐藏>", text)
    text = re.sub(r"(?i)bearer\s+[^\s\"']+", "Bearer <已隐藏>", text)
    text = re.sub(r"(?i)(?:api[_-]?key|token|authorization)[\"']?\s*[:=]\s*[\"']?[^\s,;\"']+", "<凭据已隐藏>", text)
    text = re.sub(r"(?i)data:[^,\s]+;base64,[A-Za-z0-9+/=]+", "<Base64 已隐藏>", text)
    return re.sub(r"[A-Za-z0-9+/]{256,}={0,2}", "<长数据已隐藏>", text)


def _safe_error(value: Any) -> str:
    text = _redact_log_string(str(value or "请求失败"))
    return text[:400]


def _log_value(value: Any, key: str = "", depth: int = 0) -> Any:
    if key.lower() in SENSITIVE_RESPONSE_KEYS:
        return "<已隐藏>"
    if depth >= 12:
        return "<嵌套内容已省略>"
    if isinstance(value, dict):
        return {
            _redact_log_string(str(name)[:200]): _log_value(item, str(name), depth + 1)
            for name, item in value.items()
        }
    if isinstance(value, list):
        return [_log_value(item, key, depth + 1) for item in value[:50]] + (
            [f"<其余 {len(value) - 50} 项已省略>"] if len(value) > 50 else []
        )
    if isinstance(value, str):
        if len(value) > 1200:
            return f"<长文本 {len(value)} 字符已省略>"
        return _redact_log_string(value)
    return value


def _log_response_body(stage: str, status: int, raw_body: str) -> None:
    try:
        payload = json.loads(raw_body)
    except (ValueError, TypeError):
        safe_body = _safe_error(raw_body) if raw_body else "<空响应体>"
    else:
        safe_body = json.dumps(_log_value(payload), ensure_ascii=False, separators=(",", ":"))
    if len(safe_body) > MAX_LOG_BODY:
        safe_body = f"{safe_body[:MAX_LOG_BODY]}...<后续内容已省略>"
    print(f"[Omni Flash] {stage} HTTP {status} 原始响应体（敏感值已隐藏）：{safe_body}")


async def _response_text(response: aiohttp.ClientResponse, stage: str) -> str:
    raw_body = await response.text()
    _log_response_body(stage, response.status, raw_body)
    return raw_body


def _error_code(payload: Any) -> str:
    if not isinstance(payload, dict):
        return ""
    data = payload.get("data")
    inner = data.get("data") if isinstance(data, dict) else None
    for source in (inner, data, payload):
        if not isinstance(source, dict):
            continue
        error = source.get("error")
        for value in (error.get("code") if isinstance(error, dict) else None, source.get("code")):
            if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{1,80}", value):
                return value.lower()
    return ""


def _response_error(payload: Any, status: int) -> str:
    detail = extract_error_message(payload, default="") if isinstance(payload, dict) else payload
    code = _error_code(payload)
    hint = ERROR_HINTS.get(code) or HTTP_HINTS.get(status, "视频接口请求失败")
    detail = _safe_error(detail) if detail else ""
    if detail == code or detail == hint:
        detail = ""
    suffix = f"：{detail}" if detail else ""
    code_note = f"，{code}" if code else ""
    return f"{hint}（HTTP {status}{code_note}）{suffix}"


def _task_error(payload: dict[str, Any]) -> str:
    code = _error_code(payload)
    hint = ERROR_HINTS.get(code, "视频任务生成失败")
    detail = extract_error_message(payload, default="")
    detail = _safe_error(detail) if detail else ""
    if detail == code or detail == hint:
        detail = ""
    code_note = f"（{code}）" if code else ""
    return f"{hint}{code_note}{f'：{detail}' if detail else ''}"


def _task_id(payload: Any) -> str | None:
    if not isinstance(payload, dict):
        return None
    data = payload.get("data")
    inner = data.get("data") if isinstance(data, dict) else None
    for source in (inner, data, payload):
        if isinstance(source, dict):
            for name in ("id", "task_id", "video_id"):
                value = source.get(name)
                if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{8,128}", value):
                    return value
    return None


def _task_status(payload: dict[str, Any]) -> str:
    data = payload.get("data")
    inner = data.get("data") if isinstance(data, dict) else None
    for source in (inner, data, payload):
        if isinstance(source, dict):
            for name in ("task_status", "task_state", "status", "state"):
                value = source.get(name)
                if value is not None and str(value).strip():
                    return str(value).strip().lower()
    return ""


def _video_url(payload: Any) -> str | None:
    value = extract_video_url(payload) if isinstance(payload, dict) else None
    if not value and isinstance(payload, dict):
        data = payload.get("data")
        for source in (data, payload):
            if isinstance(source, dict):
                output = source.get("output")
                if isinstance(output, dict):
                    value = output.get("video_url") or output.get("url")
                    if value:
                        break
    return value if isinstance(value, str) and len(value) <= 8192 and value.startswith(("https://", "http://")) else None


class OmniFlashClient:
    def __init__(self, *, base_url: str | None = None, api_key: str | None = None):
        self.base_url = (base_url or get_base_url_by_route()).rstrip("/")
        self.api_key = api_key or get_api_key_or_raise("O1KEY_API_KEY")

    async def _download_completed(
        self, session: aiohttp.ClientSession, task_id: str, save_path: str,
        headers: dict[str, str], status_payload: dict[str, Any],
    ) -> None:
        content_url = f"{self.base_url}/v1/videos/{task_id}/content"
        result_url = _video_url(status_payload)
        download_url, download_headers = content_url, headers
        try:
            async with session.get(content_url, headers=headers, allow_redirects=True) as response:
                if response.status >= 300:
                    raw_body = await _response_text(response, "下载")
                    if not result_url:
                        try:
                            error = json.loads(raw_body)
                        except ValueError:
                            error = raw_body
                        raise RuntimeError(_response_error(error, response.status))
                    download_url, download_headers = result_url, {}
                elif "json" in response.headers.get("Content-Type", "").lower():
                    raw_body = await _response_text(response, "下载")
                    try:
                        content_payload = json.loads(raw_body)
                    except ValueError:
                        raise RuntimeError("视频下载接口返回了无效 JSON") from None
                    if _error_code(content_payload) in ERROR_HINTS:
                        raise RuntimeError(_task_error(content_payload))
                    download_url = _video_url(content_payload) or result_url
                    if not download_url:
                        raise RuntimeError("任务已完成，但下载响应未提供视频地址")
                    download_headers = {}
                else:
                    print(f"[Omni Flash] 下载 HTTP {response.status} 响应体：<视频二进制，未打印>")
        except (aiohttp.ClientError, asyncio.TimeoutError):
            # The streaming downloader handles transient connection failures and resumes.
            if result_url:
                download_url, download_headers = result_url, {}
        await download_video_to_file(
            session, download_url, save_path, headers=download_headers or None,
            label="Omni Flash 视频",
        )

    async def generate(
        self, body: dict[str, Any], save_path: str,
        progress: Callable[[str, int, str], None] | None = None,
    ) -> str:
        headers = {"Authorization": f"Bearer {self.api_key}"}
        submit_headers = dict(headers)
        if body.get("model") == "omni_flash_abra_edit":
            submit_headers["X-No-Watermark"] = "video"
        timeout = aiohttp.ClientTimeout(total=120, connect=30)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(
                f"{self.base_url}/v1/videos", headers=submit_headers,
                **_submission_payload(body),
            ) as response:
                raw_body = await _response_text(response, "提交")
                try:
                    payload = json.loads(raw_body)
                except ValueError:
                    if response.status >= 300:
                        raise RuntimeError(_response_error(raw_body, response.status)) from None
                    raise RuntimeError("提交接口未返回有效 JSON") from None
                if response.status >= 300:
                    raise RuntimeError(_response_error(payload, response.status))
                if _error_code(payload) in ERROR_HINTS:
                    raise RuntimeError(_task_error(payload))
            task_id = _task_id(payload)
            if not task_id:
                raise RuntimeError("接口未返回有效任务 ID")
            if progress:
                progress("polling", 0, task_id)
            deadline = time.monotonic() + POLL_DEADLINE_SECONDS
            last_status = ""
            while time.monotonic() < deadline:
                await interruptible_sleep(POLL_SECONDS)
                check_interrupt()
                try:
                    async with session.get(f"{self.base_url}/v1/videos/{task_id}", headers=headers) as response:
                        raw_body = await _response_text(response, "查询")
                        if response.status in {408, 500, 502, 503, 504}:
                            continue
                        try:
                            status_payload = json.loads(raw_body)
                        except ValueError:
                            if response.status >= 300:
                                raise RuntimeError(_response_error(raw_body, response.status)) from None
                            raise
                        if response.status >= 300:
                            raise RuntimeError(_response_error(status_payload, response.status))
                except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, ValueError):
                    continue
                if not isinstance(status_payload, dict):
                    continue
                status = _task_status(status_payload)
                last_status = status or last_status
                code = _error_code(status_payload)
                if is_failure_status(status) or code in ERROR_HINTS or (
                    not status and code and code not in {"ok", "success", "0"}
                ):
                    raise RuntimeError(_task_error(status_payload))
                if is_success_status(status) or _video_url(status_payload):
                    if progress:
                        progress("downloading", 100, task_id)
                    try:
                        await self._download_completed(session, task_id, save_path, headers, status_payload)
                    except InterruptProcessingException:
                        raise
                    except Exception as exc:
                        raise RuntimeError(_safe_error(exc)) from None
                    return task_id
                if progress:
                    progress("polling", extract_progress(status_payload), task_id)
            status_note = f"，最后状态：{_safe_error(last_status)}" if last_status else ""
            raise TimeoutError(f"Omni Flash 任务 {task_id} 等待超时{status_note}")