"""Jimeng (Volcengine Visual) video generation helpers.

This module implements a minimal client for 即梦AI 视频生成 (visual.volcengineapi.com),
using Volcengine OpenAPI style HMAC-SHA256 signing with AK/SK.

Currently supported:
- Text-to-video (T2V) via req_key like 'jimeng_t2v_v30'
- Submit task + poll result + download MP4
"""

from __future__ import annotations

import hashlib
import hmac
import json
import logging
import os
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, Iterable, List, Optional, Tuple
from urllib.parse import quote

import requests


logger = logging.getLogger(__name__)


DEFAULT_JIMENG_VIDEO_MODEL = "jimeng_t2v_v30"
DEFAULT_JIMENG_ENDPOINT = "https://visual.volcengineapi.com"
DEFAULT_JIMENG_REGION = "cn-north-1"
DEFAULT_JIMENG_SERVICE = "cv"
DEFAULT_JIMENG_VERSION = "2022-08-31"


_ALLOWED_ASPECT_RATIOS = {"16:9", "4:3", "1:1", "3:4", "9:16", "21:9"}
_ALLOWED_FRAMES = {121, 241}  # 5s/10s at 24fps (+1 frame)


@dataclass(frozen=True)
class JimengAuth:
    access_key: str
    secret_key: str


def _sha256_hex(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _hmac_sha256(key: bytes, msg: str) -> bytes:
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()


def _pct_encode(s: str) -> str:
    return quote(s, safe="-_.~")


def _canonical_query(params: Iterable[Tuple[str, str]]) -> str:
    items = sorted(((k, v) for k, v in params), key=lambda kv: (kv[0], kv[1]))
    return "&".join(f"{_pct_encode(k)}={_pct_encode(v)}" for k, v in items)


def _sign_headers(
    *,
    method: str,
    host: str,
    canonical_uri: str,
    query_params: List[Tuple[str, str]],
    headers: Dict[str, str],
    body: bytes,
    auth: JimengAuth,
    region: str,
    service: str,
    request_datetime: str,
) -> Dict[str, str]:
    """Sign a request using Volcengine OpenAPI HMAC-SHA256 (header-based)."""
    algorithm = "HMAC-SHA256"
    short_date = request_datetime[:8]

    payload_hash = _sha256_hex(body)

    # Minimal set required for this service; keep consistent with working calls.
    signed_header_keys = ["content-type", "host", "x-content-sha256", "x-date"]

    lower_headers = {k.lower(): " ".join(v.strip().split()) for k, v in headers.items()}
    lower_headers["host"] = host
    lower_headers["x-date"] = request_datetime
    lower_headers["x-content-sha256"] = payload_hash

    canonical_headers = "".join(f"{k}:{lower_headers[k]}\n" for k in signed_header_keys)
    signed_headers = ";".join(signed_header_keys)

    canonical_request = (
        f"{method.upper()}\n"
        f"{canonical_uri}\n"
        f"{_canonical_query(query_params)}\n"
        f"{canonical_headers}\n"
        f"{signed_headers}\n"
        f"{payload_hash}"
    )

    credential_scope = f"{short_date}/{region}/{service}/request"
    string_to_sign = (
        f"{algorithm}\n"
        f"{request_datetime}\n"
        f"{credential_scope}\n"
        f"{_sha256_hex(canonical_request.encode('utf-8'))}"
    )

    k_date = _hmac_sha256(auth.secret_key.encode("utf-8"), short_date)
    k_region = hmac.new(k_date, region.encode("utf-8"), hashlib.sha256).digest()
    k_service = hmac.new(k_region, service.encode("utf-8"), hashlib.sha256).digest()
    k_signing = hmac.new(k_service, b"request", hashlib.sha256).digest()
    signature = hmac.new(k_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()

    authorization = (
        f"{algorithm} Credential={auth.access_key}/{credential_scope}, "
        f"SignedHeaders={signed_headers}, "
        f"Signature={signature}"
    )

    out = dict(headers)
    out["X-Date"] = request_datetime
    out["X-Content-Sha256"] = payload_hash
    out["Authorization"] = authorization
    return out


def _resolve_auth(
    access_key: Optional[str] = None,
    secret_key: Optional[str] = None,
) -> JimengAuth:
    access_key = access_key or os.environ.get("VOLC_ACCESSKEY")
    secret_key = secret_key or os.environ.get("VOLC_SECRETKEY")
    if not access_key or not secret_key:
        raise ValueError("Missing Jimeng credentials: set VOLC_ACCESSKEY and VOLC_SECRETKEY (or pass explicitly).")
    return JimengAuth(access_key=access_key, secret_key=secret_key)


class JimengVideoClient:
    """Minimal Jimeng client for async video generation tasks."""

    def __init__(
        self,
        *,
        access_key: Optional[str] = None,
        secret_key: Optional[str] = None,
        endpoint: str = DEFAULT_JIMENG_ENDPOINT,
        region: str = DEFAULT_JIMENG_REGION,
        service: str = DEFAULT_JIMENG_SERVICE,
        version: str = DEFAULT_JIMENG_VERSION,
        session: Optional[requests.Session] = None,
    ) -> None:
        self._auth = _resolve_auth(access_key, secret_key)
        self._endpoint = (endpoint or DEFAULT_JIMENG_ENDPOINT).rstrip("/")
        self._region = region or DEFAULT_JIMENG_REGION
        self._service = service or DEFAULT_JIMENG_SERVICE
        self._version = version or DEFAULT_JIMENG_VERSION
        self._session = session or requests.Session()

    def _call(self, *, action: str, body_obj: Dict[str, Any], timeout_s: float = 30.0) -> Dict[str, Any]:
        url = f"{self._endpoint}/"
        host = self._endpoint.split("://", 1)[-1].split("/", 1)[0]
        query_params = [("Action", action), ("Version", self._version)]
        body = json.dumps(body_obj, ensure_ascii=False, separators=(",", ":")).encode("utf-8")

        request_datetime = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        base_headers = {
            "Content-Type": "application/json",
            "Region": self._region,
            "Service": self._service,
        }
        signed_headers = _sign_headers(
            method="POST",
            host=host,
            canonical_uri="/",
            query_params=query_params,
            headers=base_headers,
            body=body,
            auth=self._auth,
            region=self._region,
            service=self._service,
            request_datetime=request_datetime,
        )

        resp = self._session.post(url, params=query_params, headers=signed_headers, data=body, timeout=timeout_s)
        # Service returns JSON even on non-200 for many cases; fail fast on non-JSON.
        payload = resp.json()
        payload["_http_status"] = resp.status_code
        return payload

    def submit_t2v(
        self,
        *,
        prompt: str,
        model: str = DEFAULT_JIMENG_VIDEO_MODEL,
        seed: int = -1,
        frames: int = 121,
        aspect_ratio: str = "16:9",
        timeout_s: float = 30.0,
    ) -> str:
        if not prompt or not isinstance(prompt, str):
            raise ValueError("prompt must be a non-empty string")
        if frames not in _ALLOWED_FRAMES:
            raise ValueError(f"frames must be one of {sorted(_ALLOWED_FRAMES)}")
        if aspect_ratio not in _ALLOWED_ASPECT_RATIOS:
            raise ValueError(f"aspect_ratio must be one of {sorted(_ALLOWED_ASPECT_RATIOS)}")

        payload = self._call(
            action="CVSync2AsyncSubmitTask",
            body_obj={
                "req_key": model,
                "prompt": prompt,
                "seed": int(seed),
                "frames": int(frames),
                "aspect_ratio": aspect_ratio,
            },
            timeout_s=timeout_s,
        )

        if payload.get("_http_status") != 200 or payload.get("code") != 10000:
            raise RuntimeError(f"Jimeng submit failed: http={payload.get('_http_status')} code={payload.get('code')} message={payload.get('message')}")

        data = payload.get("data") or {}
        task_id = data.get("task_id")
        if not task_id:
            raise RuntimeError("Jimeng submit returned no task_id")
        return str(task_id)

    def get_result(self, *, model: str, task_id: str, timeout_s: float = 30.0) -> Dict[str, Any]:
        payload = self._call(
            action="CVSync2AsyncGetResult",
            body_obj={"req_key": model, "task_id": str(task_id)},
            timeout_s=timeout_s,
        )
        if payload.get("_http_status") != 200 or payload.get("code") != 10000:
            raise RuntimeError(f"Jimeng get-result failed: http={payload.get('_http_status')} code={payload.get('code')} message={payload.get('message')}")
        data = payload.get("data") or {}
        if not isinstance(data, dict):
            raise RuntimeError("Jimeng get-result returned non-dict data")
        return data

    def wait_for_video_url(
        self,
        *,
        model: str,
        task_id: str,
        poll_interval_s: float = 5.0,
        timeout_s: float = 600.0,
    ) -> str:
        if poll_interval_s <= 0:
            raise ValueError("poll_interval_s must be > 0")
        if timeout_s <= 0:
            raise ValueError("timeout_s must be > 0")

        start = time.time()
        last_status: Optional[str] = None
        while True:
            data = self.get_result(model=model, task_id=task_id)
            status = str(data.get("status") or "")
            if status and status != last_status:
                logger.info("Jimeng task status=%s task_id=%s", status, task_id)
                last_status = status

            if status == "done":
                video_url = str((data.get("video_url") or "")).strip()
                if not video_url:
                    raise RuntimeError("Jimeng task done but missing video_url")
                return video_url
            if status in ("expired", "not_found"):
                raise RuntimeError(f"Jimeng task terminal status={status}")

            if (time.time() - start) >= timeout_s:
                raise TimeoutError(f"Jimeng video generation timed out after {timeout_s:.0f}s (status={status or 'unknown'})")
            time.sleep(poll_interval_s)

    def download_video_bytes(self, *, video_url: str, timeout_s: float = 120.0) -> bytes:
        if not video_url:
            raise ValueError("video_url is required")
        resp = self._session.get(video_url, timeout=timeout_s)
        resp.raise_for_status()
        return resp.content

