import asyncio
import os
import time
from typing import Any, Callable, Dict, Optional

import aiohttp


# 视频任务轮询总时长上限（秒）。超过后停止等待并抛出 TimeoutError，
# 避免任务在服务端长时间无进展时无限轮询、迫使用户手动中断。
POLL_DEADLINE_SECONDS = 2000


class PollDeadline:
    """轮询看门狗：累计等待超过上限即抛出 TimeoutError。

    用法：
        deadline = PollDeadline(label="K3 图生视频")
        while True:
            deadline.check()
            ...
    """

    def __init__(self, seconds: float = POLL_DEADLINE_SECONDS, label: str = "视频任务"):
        self.seconds = seconds
        self.label = label
        self.start = time.time()

    def elapsed(self) -> float:
        return time.time() - self.start

    def check(self) -> None:
        if self.elapsed() >= self.seconds:
            raise TimeoutError(
                f"{self.label} 轮询已超过 {self.seconds:.0f}s 仍未完成，已停止等待。"
                f"任务可能仍在服务端生成，请稍后重试。"
            )


try:
    from comfy.model_management import processing_interrupted, InterruptProcessingException
    INTERRUPT_AVAILABLE = True
except Exception:
    INTERRUPT_AVAILABLE = False
    processing_interrupted = lambda: False

    class InterruptProcessingException(Exception):
        pass


SUCCESS_STATUSES = {"succeed", "succeeded", "success", "completed", "done", "finished"}
FAILURE_STATUSES = {
    "fail",
    "failed",
    "failure",
    "error",
    "expired",
    "timeout",
    "timed_out",
    "cancel",
    "canceled",
    "cancelled",
    "rejected",
}


SEEDANCE_COPYRIGHT_RESTRICTION_ERRORS = (
    "The request failed because the output video may be related to copyright restriction",
    "The request failed because the output video may be related to copyright restrictions",
)
SEEDANCE_COPYRIGHT_RESTRICTION_MESSAGE = "输出视频触发版权审查被拒绝生成！"


def format_seedance_generation_error(value: Any) -> str:
    """Normalize errors shared by the two interactive Seedance video nodes."""
    message = str(value or "视频生成失败")
    normalized_message = message.lower()
    if any(
        marker.lower() in normalized_message
        for marker in SEEDANCE_COPYRIGHT_RESTRICTION_ERRORS
    ):
        return SEEDANCE_COPYRIGHT_RESTRICTION_MESSAGE
    return message


def check_interrupt() -> None:
    if INTERRUPT_AVAILABLE and processing_interrupted():
        raise InterruptProcessingException()


async def interruptible_sleep(seconds: float, step: float = 0.2) -> None:
    elapsed = 0.0
    while elapsed < seconds:
        check_interrupt()
        delay = min(step, seconds - elapsed)
        await asyncio.sleep(delay)
        elapsed += delay
    check_interrupt()


async def run_with_interrupt(coro, step: float = 0.2):
    task = asyncio.ensure_future(coro)
    try:
        while not task.done():
            check_interrupt()
            await asyncio.wait({task}, timeout=step)
        check_interrupt()
        return await task
    except InterruptProcessingException:
        task.cancel()
        try:
            await task
        except BaseException:
            pass
        raise


def _as_dict(value: Any) -> Dict[str, Any]:
    return value if isinstance(value, dict) else {}


def _nested_payloads(payload: Dict[str, Any]):
    root = _as_dict(payload)
    data = _as_dict(root.get("data"))
    inner = _as_dict(data.get("data"))
    return root, data, inner


def extract_status(payload: Dict[str, Any]) -> str:
    root, data, inner = _nested_payloads(payload)
    keys = ("status", "task_status", "state", "task_state")
    statuses = []
    for source in (data, inner, root):
        for key in keys:
            value = source.get(key)
            if value is not None and str(value).strip():
                statuses.append(str(value).strip().lower())
    for status in statuses:
        if status in FAILURE_STATUSES or any(
            token in status for token in ("fail", "error", "reject", "timeout", "cancel")
        ):
            return status
    for status in statuses:
        if status in SUCCESS_STATUSES:
            return status
    return statuses[0] if statuses else ""


def extract_progress(payload: Dict[str, Any]) -> int:
    root, data, inner = _nested_payloads(payload)
    for source in (data, inner, root):
        value = source.get("progress")
        if value is None:
            continue
        try:
            return max(0, min(100, int(float(str(value).strip().rstrip("%")))))
        except (TypeError, ValueError):
            return 0
    return 0


def extract_error_message(payload: Dict[str, Any], default: str = "未知错误") -> str:
    root, data, inner = _nested_payloads(payload)
    keys = (
        "fail_reason",
        "failure_reason",
        "task_status_msg",
        "status_msg",
        "error_message",
        "message",
        "msg",
        "reason",
        "detail",
        "details",
    )
    for source in (data, inner, root):
        error = source.get("error")
        if isinstance(error, dict):
            for key in ("message", "msg", "detail", "reason", "code"):
                value = error.get(key)
                if value:
                    return str(value)
        elif error:
            return str(error)

        for key in keys:
            value = source.get(key)
            if value:
                return str(value)
    return default


def extract_video_url(payload: Dict[str, Any]) -> str | None:
    root, data, inner = _nested_payloads(payload)
    for source in (data, inner, root):
        for key in ("video_url", "result_url", "url", "download_url"):
            value = source.get(key)
            if value:
                return str(value)

        result = _as_dict(source.get("result"))
        for key in ("video_url", "result_url", "url", "download_url"):
            value = result.get(key)
            if value:
                return str(value)

        content = _as_dict(source.get("content"))
        value = content.get("video_url") or content.get("url")
        if value:
            return str(value)

        task_result = _as_dict(source.get("task_result"))
        videos = task_result.get("videos")
        if isinstance(videos, list) and videos:
            first = _as_dict(videos[0])
            value = first.get("url") or first.get("video_url")
            if value:
                return str(value)

        # 腾讯 Kling（v3-t）渠道：完成时视频地址在 metadata.url
        metadata = _as_dict(source.get("metadata"))
        value = metadata.get("url") or metadata.get("video_url")
        if value:
            return str(value)
    return None


def is_success_status(status: str) -> bool:
    return status in SUCCESS_STATUSES


def is_failure_status(status: str, payload: Dict[str, Any] | None = None) -> bool:
    if status in FAILURE_STATUSES:
        return True
    if any(token in status for token in ("fail", "error", "reject", "timeout", "cancel")):
        return True
    if payload is None:
        return False
    root, data, inner = _nested_payloads(payload)
    failure_keys = ("error", "fail_reason", "failure_reason", "task_status_msg", "error_message")
    return any(any(source.get(key) for key in failure_keys) for source in (data, inner, root))


# ═══════════════════════════════════════════════════════════════════════════════
# 视频下载：抗超时 / 可断点续传 / 无限重试 / 可随时取消
# ═══════════════════════════════════════════════════════════════════════════════
#
# 设计目标（针对“视频已在服务端生成成功、后台已扣费，必须把成品拿到手”的场景）：
#   1. 不用固定 total 超时一刀切大文件——只要持续有数据就一直下载；
#      用 sock_read 检测“卡死”（连续 N 秒收不到任何字节）才判定异常。
#   2. 网络抖动 / 超时 / 5xx / 连接中断 → 退避后无限重试，直到成功。
#   3. 已下载的字节用 HTTP Range 断点续传，不从 0 重来。
#   4. 永久性错误（403/404/410 等）快速失败，不做无意义的死循环。
#   5. 全程可被 ComfyUI 的“取消”随时打断（每个 attempt 包在 run_with_interrupt 中，
#      分块写入与退避等待都会检查中断）。

# 连续多少秒收不到任何数据就判定当前连接卡死（触发重试，而非整体失败）
DOWNLOAD_SOCK_READ_TIMEOUT = 120
# 建立连接的超时
DOWNLOAD_CONNECT_TIMEOUT = 30
# 重试退避：起始 / 上限（秒）
DOWNLOAD_RETRY_BASE_DELAY = 2.0
DOWNLOAD_RETRY_MAX_DELAY = 30.0
# 视为“永久失败、无需重试”的 HTTP 状态码
DOWNLOAD_PERMANENT_STATUS = {400, 401, 403, 404, 405, 410, 451}


class _PermanentDownloadError(RuntimeError):
    """不可重试的下载错误（如 403/404）。"""


class _IncompleteDownloadError(RuntimeError):
    """连接被提前关闭、文件未下完，需要续传重试。"""


def _download_timeout(sock_read: float) -> aiohttp.ClientTimeout:
    # total=None：不限制总时长，让缓慢但持续的大文件下载得以完成；
    # sock_connect/sock_read：分别约束“连接建立”和“两次收包之间”的最大间隔。
    return aiohttp.ClientTimeout(
        total=None,
        connect=None,
        sock_connect=DOWNLOAD_CONNECT_TIMEOUT,
        sock_read=sock_read,
    )


async def _stream_once(
    session: aiohttp.ClientSession,
    url: str,
    fobj,
    *,
    headers: Optional[dict],
    resume_from: int,
    sock_read: float,
    chunk_size: int,
    on_bytes: Optional[Callable[[int], None]],
) -> int:
    """发起一次 GET 并把响应体写入已打开的文件对象 fobj。

    返回“本次结束后文件应当达到的总字节数”（已知时），未知时返回 -1。
    若服务端支持 Range 且 resume_from>0，则带上 Range 头从断点继续；
    否则从头下载（必要时先 truncate）。出错时抛异常，由上层决定是否重试。
    """
    req_headers = dict(headers) if headers else {}
    if resume_from > 0:
        req_headers["Range"] = f"bytes={resume_from}-"

    timeout = _download_timeout(sock_read)
    async with session.get(
        url, headers=req_headers or None, timeout=timeout, allow_redirects=True
    ) as resp:
        status = resp.status

        # 206：服务端接受断点续传，从 resume_from 续写。
        # 200：服务端忽略 Range（或本就从头下），需从文件起点重写。
        base = resume_from
        if resume_from > 0 and status == 200:
            fobj.seek(0)
            fobj.truncate(0)
            base = 0
        elif status not in (200, 206):
            text = ""
            try:
                text = (await resp.text())[:500]
            except Exception:
                pass
            if status in DOWNLOAD_PERMANENT_STATUS:
                raise _PermanentDownloadError(f"视频下载失败 ({status})：{text}")
            raise RuntimeError(f"视频下载失败 ({status})：{text}")

        # 解析“完整文件总大小”，用于检测连接被提前关闭导致的截断。
        expected_total = _expected_total_size(resp, base)

        async for chunk in resp.content.iter_chunked(chunk_size):
            check_interrupt()
            if chunk:
                fobj.write(chunk)
                if on_bytes:
                    on_bytes(len(chunk))
        fobj.flush()
        return expected_total


def _expected_total_size(resp: aiohttp.ClientResponse, base: int) -> int:
    """根据响应头推断完整文件总字节数；无法判断时返回 -1。"""
    # Content-Range: bytes start-end/total → total 即完整大小
    cr = resp.headers.get("Content-Range", "")
    if "/" in cr:
        tail = cr.rsplit("/", 1)[-1].strip()
        if tail.isdigit():
            return int(tail)
    # Content-Length 是“本次响应体长度”，加上已续传的 base 即完整大小
    cl = resp.headers.get("Content-Length")
    if cl is not None and cl.isdigit():
        return base + int(cl)
    return -1


async def download_video_to_file(
    session: aiohttp.ClientSession,
    url: str,
    save_path: str,
    *,
    headers: Optional[dict] = None,
    sock_read: float = DOWNLOAD_SOCK_READ_TIMEOUT,
    chunk_size: int = 1024 * 1024,
    label: str = "视频",
    on_bytes: Optional[Callable[[int], None]] = None,
    max_retries: Optional[int] = None,
) -> str:
    """把远程视频下载到 save_path，抗超时 + 断点续传 + 无限重试 + 可取消。

    - 已生成成功的视频务必拿到手：默认 max_retries=None 表示对“可重试错误”
      （网络中断 / 超时 / 5xx / 连接失败）一直重试，直到成功。
    - 永久性错误（403/404/410 等）立即抛出，不做无意义重试。
    - 通过 HTTP Range 从已落盘的字节处续传，不重复下载。
    - 通过 check_interrupt() 全程响应 ComfyUI 取消。

    返回 save_path。
    """
    parent = os.path.dirname(save_path)
    if parent:
        os.makedirs(parent, exist_ok=True)

    # 起始先清空目标文件：避免对调用方残留的旧文件做错误续传（续传只针对本次调用
    # 内部已写入的字节）。
    try:
        with open(save_path, "wb"):
            pass
    except OSError:
        pass

    attempt = 0
    while True:
        check_interrupt()
        # 仅对“本次调用已落盘”的字节做断点续传。
        resume_from = 0
        if os.path.isfile(save_path):
            try:
                resume_from = os.path.getsize(save_path)
            except OSError:
                resume_from = 0

        # 已有部分文件 → 追加续写；否则新建。
        mode = "r+b" if resume_from > 0 else "wb"
        try:
            with open(save_path, mode) as f:
                if resume_from > 0:
                    f.seek(0, os.SEEK_END)
                expected_total = await run_with_interrupt(
                    _stream_once(
                        session, url, f,
                        headers=headers,
                        resume_from=resume_from,
                        sock_read=sock_read,
                        chunk_size=chunk_size,
                        on_bytes=on_bytes,
                    )
                )
            # 走到这里说明本次 GET 的响应体已读尽、连接已关闭。
            size = os.path.getsize(save_path) if os.path.isfile(save_path) else 0
            if size <= 0:
                raise RuntimeError(f"{label}下载失败：保存后的文件为空。")
            # 连接被提前关闭（截断）：实际大小 < 服务端声明的完整大小 → 续传重试。
            if expected_total > 0 and size < expected_total:
                raise _IncompleteDownloadError(
                    f"{label}下载不完整：{size}/{expected_total} 字节，将续传。"
                )
            return save_path

        except InterruptProcessingException:
            raise
        except _PermanentDownloadError:
            raise
        except (_IncompleteDownloadError, aiohttp.ClientError, asyncio.TimeoutError, OSError) as e:
            attempt += 1
            if max_retries is not None and attempt > max_retries:
                raise RuntimeError(f"{label}下载失败（已重试 {max_retries} 次）：{e}") from None
            delay = min(DOWNLOAD_RETRY_BASE_DELAY * (2 ** (attempt - 1)), DOWNLOAD_RETRY_MAX_DELAY)
            done = "续传" if (os.path.isfile(save_path) and os.path.getsize(save_path) > 0) else "重连"
            print(f"[{label}] 下载中断（{type(e).__name__}），{delay:.0f}s 后{done}重试（第 {attempt} 次）...")
            await interruptible_sleep(delay)


async def download_video_bytes(
    session: aiohttp.ClientSession,
    url: str,
    *,
    headers: Optional[dict] = None,
    label: str = "视频",
    **kwargs,
) -> bytes:
    """与 download_video_to_file 相同的健壮性，但返回内存中的 bytes。

    内部仍落盘到临时文件以支持断点续传，读出后删除。
    """
    import tempfile

    fd, tmp_path = tempfile.mkstemp(suffix=".bin", prefix="o1key_dl_")
    os.close(fd)
    try:
        await download_video_to_file(
            session, url, tmp_path, headers=headers, label=label, **kwargs
        )
        with open(tmp_path, "rb") as f:
            return f.read()
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass