"""Jimeng (Volcengine Visual) image 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-image (T2I) via req_key like 'jimeng_t2i_v40'
- Submit async task + poll result + download image bytes
"""

from __future__ import annotations

import base64
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_IMAGE_MODEL = "jimeng_t2i_v40"
DEFAULT_JIMENG_ENDPOINT = "https://visual.volcengineapi.com"
DEFAULT_JIMENG_REGION = "cn-north-1"
DEFAULT_JIMENG_SERVICE = "cv"
DEFAULT_JIMENG_VERSION = "2022-08-31"


@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)

    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)


def _extract_image_urls(data: Dict[str, Any]) -> List[str]:
    urls: List[str] = []
    raw = data.get("image_urls")
    if isinstance(raw, list):
        urls.extend([str(x) for x in raw if str(x).strip()])
    raw_one = data.get("image_url")
    if raw_one:
        urls.append(str(raw_one))
    # Some variants may nest results.
    nested = data.get("result")
    if isinstance(nested, str) and nested.strip().startswith(("{", "[")):
        try:
            nested_obj = json.loads(nested)
            if isinstance(nested_obj, dict):
                urls.extend(_extract_image_urls(nested_obj))
        except Exception:
            pass
    if isinstance(nested, dict):
        urls.extend(_extract_image_urls(nested))
    # De-dup preserving order.
    out: List[str] = []
    seen: set[str] = set()
    for u in urls:
        if u not in seen:
            out.append(u)
            seen.add(u)
    return out


def _extract_image_base64(data: Dict[str, Any]) -> Optional[bytes]:
    raw = data.get("binary_data_base64")
    if isinstance(raw, str) and raw.strip():
        try:
            return base64.b64decode(raw, validate=False)
        except Exception:
            return None
    if isinstance(raw, list) and raw:
        first = raw[0]
        if isinstance(first, str) and first.strip():
            try:
                return base64.b64decode(first, validate=False)
            except Exception:
                return None
    return None


class JimengImageClient:
    """Minimal Jimeng client for async image 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)
        payload = resp.json()
        payload["_http_status"] = resp.status_code
        return payload

    def submit_t2i(
        self,
        *,
        prompt: str,
        model: str = DEFAULT_JIMENG_IMAGE_MODEL,
        seed: int = -1,
        width: Optional[int] = None,
        height: Optional[int] = None,
        timeout_s: float = 30.0,
        extra: Optional[Dict[str, Any]] = None,
    ) -> str:
        if not prompt or not isinstance(prompt, str):
            raise ValueError("prompt must be a non-empty string")

        body: Dict[str, Any] = {
            "req_key": model,
            "prompt": prompt,
            "seed": int(seed),
        }
        if width is not None:
            body["width"] = int(width)
        if height is not None:
            body["height"] = int(height)
        if extra:
            body.update(extra)

        payload = self._call(action="CVSync2AsyncSubmitTask", body_obj=body, 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_image_result(
        self,
        *,
        model: str,
        task_id: str,
        poll_interval_s: float = 2.0,
        timeout_s: float = 300.0,
    ) -> Dict[str, Any]:
        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":
                return data
            if status in ("expired", "not_found"):
                raise RuntimeError(f"Jimeng task terminal status={status}")

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

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

    def wait_and_download_first_image(
        self,
        *,
        model: str,
        task_id: str,
        poll_interval_s: float = 2.0,
        timeout_s: float = 300.0,
    ) -> tuple[bytes, Dict[str, Any]]:
        data = self.wait_for_image_result(
            model=model,
            task_id=task_id,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
        )
        b64 = _extract_image_base64(data)
        if b64:
            return b64, data

        urls = _extract_image_urls(data)
        if not urls:
            raise RuntimeError("Jimeng task done but missing image_urls/binary_data_base64")
        return self.download_image_bytes(url=urls[0]), data


__all__ = [
    "JimengImageClient",
    "DEFAULT_JIMENG_IMAGE_MODEL",
]

