microdao-daarion/services/aurora-service/app/kling.py

"""Kling AI API client for video generation and enhancement."""
from __future__ import annotations

import hashlib
import hmac
import json
import logging
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import urllib.request
import urllib.error

logger = logging.getLogger(__name__)

KLING_ACCESS_KEY = os.getenv("KLING_ACCESS_KEY", "").strip()
KLING_SECRET_KEY = os.getenv("KLING_SECRET_KEY", "").strip()
KLING_BASE_URL = os.getenv("KLING_BASE_URL", "https://api.klingai.com")
KLING_TIMEOUT = int(os.getenv("KLING_TIMEOUT", "60"))


def _kling_sign(access_key: str, secret_key: str) -> str:
    """Generate Kling AI Bearer token via HMAC-SHA256 JWT-style signing."""
    import base64
    header = base64.urlsafe_b64encode(json.dumps({"alg": "HS256", "typ": "JWT"}).encode()).rstrip(b"=").decode()
    now = int(time.time())
    payload = base64.urlsafe_b64encode(json.dumps({
        "iss": access_key,
        "exp": now + 1800,
        "nbf": now - 5,
    }).encode()).rstrip(b"=").decode()
    msg = f"{header}.{payload}"
    sig = hmac.new(secret_key.encode(), msg.encode(), hashlib.sha256).digest()
    sig_b64 = base64.urlsafe_b64encode(sig).rstrip(b"=").decode()
    return f"{msg}.{sig_b64}"


def _kling_headers() -> Dict[str, str]:
    if not KLING_ACCESS_KEY or not KLING_SECRET_KEY:
        raise RuntimeError(
            "Kling credentials are not configured. Set KLING_ACCESS_KEY and KLING_SECRET_KEY."
        )
    token = _kling_sign(KLING_ACCESS_KEY, KLING_SECRET_KEY)
    return {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }


def _kling_request(method: str, path: str, body: Optional[Dict] = None, timeout: int = KLING_TIMEOUT) -> Dict[str, Any]:
    url = f"{KLING_BASE_URL}{path}"
    data = json.dumps(body).encode() if body else None
    req = urllib.request.Request(url, data=data, headers=_kling_headers(), method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw)
    except urllib.error.HTTPError as e:
        err_body = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"Kling API {method} {path} → {e.code}: {err_body}") from e


def _kling_request_with_fallback(
    method: str,
    paths: List[str],
    body: Optional[Dict] = None,
    timeout: int = KLING_TIMEOUT,
) -> Dict[str, Any]:
    """Try several endpoint variants to tolerate provider path drift/gateway prefixes."""
    last_error: Optional[str] = None
    tried: List[str] = []
    for path in paths:
        tried.append(path)
        try:
            return _kling_request(method, path, body=body, timeout=timeout)
        except RuntimeError as e:
            msg = str(e)
            last_error = msg
            # 404 likely means wrong endpoint path; try next candidate.
            if "→ 404:" in msg:
                continue
            # Non-404 errors are usually actionable immediately.
            raise
    raise RuntimeError(
        f"Kling API endpoint mismatch for {method}. Tried: {tried}. Last error: {last_error or 'unknown'}"
    )


# ── Video Enhancement (Video-to-Video) ──────────────────────────────────────

def kling_video_enhance(
    *,
    video_url: Optional[str] = None,
    video_id: Optional[str] = None,
    prompt: str = "",
    negative_prompt: str = "noise, blur, artifacts",
    mode: str = "pro",
    duration: str = "5",
    cfg_scale: float = 0.5,
    callback_url: Optional[str] = None,
) -> Dict[str, Any]:
    """Submit a video-to-video enhancement task to Kling AI.

    Args:
        video_url: Direct URL to input video.
        video_id: Kling resource ID for previously uploaded video.
        prompt: Text guidance for enhancement.
        negative_prompt: Things to avoid.
        mode: 'std' or 'pro'.
        duration: '5' or '10' seconds.
        cfg_scale: 0.0-1.0, how strongly to follow prompt.
        callback_url: Webhook for completion notification.

    Returns:
        Task response dict with task_id.
    """
    if not video_url and not video_id:
        raise ValueError("Either video_url or video_id must be provided")

    payload: Dict[str, Any] = {
        "model": f"kling-v1-5",
        "mode": mode,
        "duration": duration,
        "cfg_scale": cfg_scale,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
    }
    if video_url:
        payload["video_url"] = video_url
    if video_id:
        payload["video_id"] = video_id
    if callback_url:
        payload["callback_url"] = callback_url

    return _kling_request_with_fallback(
        "POST",
        ["/v1/videos/video2video", "/kling/v1/videos/video2video"],
        body=payload,
    )


def kling_video_generate(
    *,
    image_b64: Optional[str] = None,
    image_url: Optional[str] = None,
    image_id: Optional[str] = None,
    prompt: str,
    negative_prompt: str = "noise, blur, artifacts, distortion",
    model: str = "kling-v1-5",
    mode: str = "pro",
    duration: str = "5",
    cfg_scale: float = 0.5,
    aspect_ratio: str = "16:9",
    callback_url: Optional[str] = None,
) -> Dict[str, Any]:
    """Generate video from image + prompt (image-to-video).

    Args:
        image_url: Source still frame URL.
        image_id: Kling resource ID for previously uploaded image.
        prompt: Animation guidance.
        model: 'kling-v1', 'kling-v1-5', 'kling-v1-6'.
        mode: 'std' or 'pro'.
        duration: '5' or '10'.
        aspect_ratio: '16:9', '9:16', '1:1'.
    """
    if not image_b64 and not image_url and not image_id:
        raise ValueError("One of image_b64 / image_url / image_id must be provided")

    payload: Dict[str, Any] = {
        "model": model,
        "mode": mode,
        "duration": duration,
        "cfg_scale": cfg_scale,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "aspect_ratio": aspect_ratio,
    }
    # Current Kling endpoint expects "image" as base64 payload string.
    # Keep url/id compatibility as a best-effort fallback for older gateways.
    if image_b64:
        payload["image"] = image_b64
    elif image_url:
        payload["image"] = image_url
    elif image_id:
        payload["image"] = image_id
    if callback_url:
        payload["callback_url"] = callback_url

    return _kling_request_with_fallback(
        "POST",
        ["/v1/videos/image2video", "/kling/v1/videos/image2video"],
        body=payload,
    )


def kling_video_generate_from_file(
    *,
    image_path: Path,
    prompt: str,
    negative_prompt: str = "noise, blur, artifacts, distortion",
    model: str = "kling-v1-5",
    mode: str = "pro",
    duration: str = "5",
    cfg_scale: float = 0.5,
    aspect_ratio: str = "16:9",
    callback_url: Optional[str] = None,
) -> Dict[str, Any]:
    """Generate video from a local image file by sending base64 payload."""
    import base64

    with image_path.open("rb") as fh:
        image_b64 = base64.b64encode(fh.read()).decode()

    return kling_video_generate(
        image_b64=image_b64,
        prompt=prompt,
        negative_prompt=negative_prompt,
        model=model,
        mode=mode,
        duration=duration,
        cfg_scale=cfg_scale,
        aspect_ratio=aspect_ratio,
        callback_url=callback_url,
    )


def kling_task_status(task_id: str) -> Dict[str, Any]:
    """Get status of any Kling task by ID."""
    return _kling_request_with_fallback(
        "GET",
        [f"/v1/tasks/{task_id}", f"/kling/v1/tasks/{task_id}"],
    )


def kling_video_task_status(task_id: str, endpoint: str = "video2video") -> Dict[str, Any]:
    """Get status of a video task."""
    return _kling_request_with_fallback(
        "GET",
        [f"/v1/videos/{endpoint}/{task_id}", f"/kling/v1/videos/{endpoint}/{task_id}"],
    )


def kling_list_models() -> Dict[str, Any]:
    """List available Kling models."""
    return _kling_request_with_fallback(
        "GET",
        ["/v1/models", "/kling/v1/models"],
    )


def kling_upload_file(file_path: Path) -> Dict[str, Any]:
    """Upload a local file to Kling storage and return resource_id."""
    import base64
    with open(file_path, "rb") as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    suffix = file_path.suffix.lstrip(".").lower()
    mime_map = {
        "mp4": "video/mp4", "mov": "video/quicktime", "avi": "video/x-msvideo",
        "jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png",
    }
    mime = mime_map.get(suffix, "application/octet-stream")
    payload = {
        "file": b64,
        "file_name": file_path.name,
        "content_type": mime,
    }
    return _kling_request_with_fallback(
        "POST",
        ["/v1/files/upload", "/v1/files", "/kling/v1/files/upload", "/kling/v1/files"],
        body=payload,
        timeout=120,
    )


def kling_poll_until_done(
    task_id: str,
    endpoint: str = "video2video",
    max_wait_sec: int = 600,
    poll_interval: int = 5,
) -> Dict[str, Any]:
    """Poll Kling task until completed/failed or timeout."""
    start = time.time()
    while True:
        status_resp = kling_video_task_status(task_id, endpoint)
        task = status_resp.get("data", {})
        state = task.get("task_status") or task.get("status") or "processing"

        if state in ("succeed", "completed", "failed", "error"):
            return status_resp

        elapsed = time.time() - start
        if elapsed >= max_wait_sec:
            raise TimeoutError(f"Kling task {task_id} timed out after {max_wait_sec}s (last status: {state})")

        logger.debug("Kling task %s status=%s elapsed=%.0fs", task_id, state, elapsed)
        time.sleep(poll_interval)


def kling_health_check() -> Dict[str, Any]:
    """Quick connectivity check — returns status dict."""
    try:
        # `/v1/models` may be disabled in some accounts/regions.
        # `/v1/videos/image2video` reliably returns code=0 when auth+endpoint are valid.
        resp = _kling_request("GET", "/v1/videos/image2video", timeout=10)
        code = resp.get("code") if isinstance(resp, dict) else None
        if code not in (None, 0, "0"):
            return {"ok": False, "error": f"Kling probe returned non-zero code: {code}", "probe": resp}
        return {"ok": True, "probe_path": "/v1/videos/image2video", "probe": resp}
    except Exception as exc:
        return {"ok": False, "error": str(exc)}