#!/usr/bin/env python3
"""
Kling v2.6 Direct API client for video generation.

Uses JWT auth against Kling's Singapore endpoint.
No third-party proxy (Pollo, LaoZhang) required.

Docs: https://docs.qingque.cn/d/home/eZQB2yXi5bHNbo0fSqJchEMDnkh
"""

import base64
import json
import os
import time
import requests
from pathlib import Path
from typing import Optional

try:
    import jwt
except ImportError:
    raise ImportError("pyjwt is required: pip install pyjwt")


DEFAULT_API_BASE = "https://api-singapore.klingai.com"


def _generate_jwt_token(access_key: str, secret_key: str) -> str:
    """Generate JWT token for Kling API authentication."""
    payload = {
        "iss": access_key,
        "exp": int(time.time()) + 1800,  # 30 min validity
        "nbf": int(time.time()) - 5
    }
    return jwt.encode(payload, secret_key, algorithm="HS256")


def _image_to_base64(image_path: str) -> str:
    """Convert image file to base64 (no data URI prefix)."""
    with open(image_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8')


def call_kling_direct(
    prompt: str,
    first_frame_path: str,
    access_key: str = None,
    secret_key: str = None,
    last_frame_path: Optional[str] = None,
    duration: int = 5,
    generate_audio: bool = True,
    api_base: str = None
) -> dict:
    """
    Submit video generation task to Kling Direct API.

    Args:
        prompt: Video generation prompt (max 2500 chars)
        first_frame_path: Path to first frame image
        access_key: Kling API access key (or KLING_ACCESS_KEY env)
        secret_key: Kling API secret key (or KLING_SECRET_KEY env)
        last_frame_path: Optional path to last frame (interpolation)
        duration: Duration in seconds (5 or 10)
        generate_audio: Generate audio track
        api_base: API base URL

    Returns:
        {"task_id": str}
    """
    access_key = access_key or os.environ.get("KLING_ACCESS_KEY")
    secret_key = secret_key or os.environ.get("KLING_SECRET_KEY")
    api_base = api_base or os.environ.get("KLING_API_BASE", DEFAULT_API_BASE)

    if not access_key or not secret_key:
        raise RuntimeError("KLING_ACCESS_KEY and KLING_SECRET_KEY must be set")

    token = _generate_jwt_token(access_key, secret_key)

    has_end_frame = last_frame_path and Path(last_frame_path).exists()

    # Kling Direct API constraint: image_tail + sound=on is not allowed.
    # Prefer frame interpolation over audio when both frames provided.
    if has_end_frame and generate_audio:
        generate_audio = False
        print(f"  Kling Direct: disabling audio (incompatible with end frame)")

    payload = {
        "model_name": "kling-v2-6",
        "mode": "pro",
        "image": _image_to_base64(first_frame_path),
        "prompt": prompt[:2500],
        "duration": str(duration),
        "sound": "on" if generate_audio else "off",
    }

    if has_end_frame:
        payload["image_tail"] = _image_to_base64(last_frame_path)
        print(f"  Kling Direct: added end frame")

    print(f"  Kling Direct: submitting task (duration={duration}s, audio={'on' if generate_audio else 'off'})...")

    response = requests.post(
        f"{api_base}/v1/videos/image2video",
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        },
        json=payload,
        timeout=60
    )

    if response.status_code != 200:
        raise RuntimeError(f"Kling Direct API error: {response.status_code} - {response.text[:500]}")

    result = response.json()

    if result.get("code") != 0:
        raise RuntimeError(f"Kling Direct error: {result.get('message', result)}")

    task_id = result.get("data", {}).get("task_id")
    if not task_id:
        raise RuntimeError(f"Kling Direct: no task_id in response: {result}")

    print(f"  Kling Direct task: {task_id}")
    return {"task_id": task_id}


def poll_kling_direct(
    task_id: str,
    access_key: str = None,
    secret_key: str = None,
    poll_interval: int = 10,
    timeout: int = 600,
    api_base: str = None
) -> dict:
    """
    Poll Kling Direct API for task completion.

    Args:
        task_id: Task ID from call_kling_direct
        access_key: Kling API access key (or KLING_ACCESS_KEY env)
        secret_key: Kling API secret key (or KLING_SECRET_KEY env)
        poll_interval: Seconds between polls
        timeout: Max wait time in seconds
        api_base: API base URL

    Returns:
        {"uri": str, "state": "COMPLETE", "duration": float}
    """
    access_key = access_key or os.environ.get("KLING_ACCESS_KEY")
    secret_key = secret_key or os.environ.get("KLING_SECRET_KEY")
    api_base = api_base or os.environ.get("KLING_API_BASE", DEFAULT_API_BASE)

    if not access_key or not secret_key:
        raise RuntimeError("KLING_ACCESS_KEY and KLING_SECRET_KEY must be set")

    status_url = f"{api_base}/v1/videos/image2video/{task_id}"
    start_time = time.time()
    consecutive_errors = 0

    while time.time() - start_time < timeout:
        # Refresh JWT token each poll (token valid 30min, but cheap to regenerate)
        token = _generate_jwt_token(access_key, secret_key)
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        try:
            response = requests.get(status_url, headers=headers, timeout=30)
        except Exception as e:
            consecutive_errors += 1
            if consecutive_errors >= 5:
                raise RuntimeError(f"Kling Direct: {consecutive_errors} consecutive network errors: {e}")
            elapsed = int(time.time() - start_time)
            print(f"  Kling Direct: network error ({elapsed}s elapsed), retrying...")
            time.sleep(poll_interval)
            continue

        consecutive_errors = 0

        if response.status_code != 200:
            raise RuntimeError(f"Kling Direct poll error: {response.status_code} - {response.text[:500]}")

        result = response.json()

        if result.get("code") != 0:
            raise RuntimeError(f"Kling Direct poll error: {result.get('message', result)}")

        data = result.get("data", {})
        task_status = data.get("task_status")
        elapsed = int(time.time() - start_time)

        print(f"  Kling Direct generating... ({elapsed}s elapsed, status: {task_status})")

        if task_status == "succeed":
            task_result = data.get("task_result", {})
            videos = task_result.get("videos", [])
            if not videos:
                raise RuntimeError("Kling Direct: no videos in completed response")

            video_info = videos[0]
            video_url = video_info.get("url")
            if not video_url:
                raise RuntimeError(f"Kling Direct: no URL in video result: {video_info}")

            print(f"  Kling Direct video ready!")
            return {
                "uri": video_url,
                "state": "COMPLETE",
                "duration": video_info.get("duration")
            }

        elif task_status == "failed":
            error_msg = data.get("task_status_msg", "Generation failed")
            raise RuntimeError(f"Kling Direct generation failed: {error_msg}")

        time.sleep(poll_interval)

    raise RuntimeError(f"Kling Direct: timeout after {timeout}s")


def download_video_kling_direct(video_url: str, output_path: str, max_retries: int = 3) -> str:
    """Download video from Kling CDN."""
    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            response = requests.get(video_url, stream=True, timeout=(15, 120))

            if response.status_code != 200:
                raise RuntimeError(f"Download error: HTTP {response.status_code}")

            expected_size = int(response.headers.get("content-length", 0))

            with open(output_path, 'wb') as f:
                downloaded = 0
                for chunk in response.iter_content(chunk_size=65536):
                    f.write(chunk)
                    downloaded += len(chunk)

            if expected_size > 0 and downloaded != expected_size:
                raise RuntimeError(f"Incomplete download: {downloaded}/{expected_size} bytes")

            if downloaded < 100_000:
                raise RuntimeError(f"File too small: {downloaded} bytes")

            return output_path

        except Exception as e:
            last_error = e
            if attempt < max_retries:
                wait = attempt * 5
                print(f"  Kling Direct download attempt {attempt}/{max_retries} failed: {e}, retrying in {wait}s...")
                time.sleep(wait)
            if os.path.exists(output_path):
                os.remove(output_path)

    raise RuntimeError(f"Kling Direct download failed after {max_retries} attempts: {last_error}")


if __name__ == "__main__":
    import sys
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).parent.parent / ".env")

    if len(sys.argv) < 2:
        print("Usage: python kling_direct.py test")
        sys.exit(1)

    if sys.argv[1] == "test":
        access_key = os.environ.get("KLING_ACCESS_KEY")
        secret_key = os.environ.get("KLING_SECRET_KEY")
        if not access_key or not secret_key:
            print("Error: KLING_ACCESS_KEY and KLING_SECRET_KEY must be set")
            sys.exit(1)

        print("Testing Kling Direct API auth...")
        token = _generate_jwt_token(access_key, secret_key)
        print(f"JWT token generated: {token[:20]}...")
        print("OK — token generation works. Use with a frame image to test full generation.")
