import gradio as gr
import requests
import base64
from pathlib import Path
import jwt
import time

# Kling AI API configuration (keys hardcoded as requested)
ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
API_URL = "https://api-singapore.klingai.com/v1/image-to-video"

def generate_jwt_token():
    """Generate JWT token for Kling AI API authentication."""
    payload = {
        "iat": int(time.time()),
        "exp": int(time.time()) + 1800,  # Token expires in 30 minutes
        "access_key": ACCESS_KEY_ID
    }
    token = jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")
    return token

def generate_video(image, prompt=""):
    """
    Call Kling AI API to generate a video from an input image.
    
    Args:
        image: Uploaded image file (from Gradio)
        prompt (str): Optional text prompt to guide video generation
    
    Returns:
        str: Path to the downloaded video or error message
    """
    if not image:
        return "Error: Please upload a valid image (PNG/JPEG)."
    
    # Convert image to base64
    try:
        with open(image, "rb") as img_file:
            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    except Exception as e:
        return f"Error: Failed to process image. Details: {str(e)}"
    
    headers = {
        "Authorization": f"Bearer {generate_jwt_token()}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "image": image_base64,
        "prompt": prompt,
        "duration": 5,
        "aspect_ratio": "16:9"
    }
    
    try:
        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        
        data = response.json()
        task_id = data.get("task_id") or data.get("taskId")
        if not task_id:
            return "Error: No task ID returned by the API."
        
        # Poll for task completion
        status_url = f"https://api-singapore.klingai.com/v1/predictions/{task_id}"
        for _ in range(60):  # Poll for up to 5 minutes
            status_response = requests.get(status_url, headers=headers, timeout=30)
            status_response.raise_for_status()
            status_data = status_response.json()
            if status_data.get("status") == "succeeded":
                video_url = status_data.get("video_url") or status_data.get("result", {}).get("video_url")
                if not video_url:
                    return "Error: No video URL in API response."
                # Download the video
                video_response = requests.get(video_url, timeout=30)
                video_response.raise_for_status()
                output_path = Path("output_video.mp4")
                with open(output_path, "wb") as f:
                    f.write(video_response.content)
                return str(output_path)
            elif status_data.get("status") == "failed":
                return "Error: Video generation failed."
            time.sleep(5)
        
        return "Error: Video generation timed out."
    
    except requests.exceptions.HTTPError as e:
        return f"Error: API request failed. Details: {str(e)}"
    except requests.exceptions.RequestException as e:
        return f"Error: Network issue. Details: {str(e)}"

def chatbot_interface(image, prompt):
    """
    Gradio interface for image-to-video generation.
    
    Args:
        image: Uploaded image file
        prompt (str): Optional text prompt
    
    Returns:
        Video file path or error message
    """
    return generate_video(image, prompt)

# Define Gradio interface
iface = gr.Interface(
    fn=chatbot_interface,
    inputs=[
        gr.Image(type="filepath", label="Upload Image (PNG/JPEG)"),
        gr.Textbox(lines=2, placeholder="Enter an optional prompt (e.g., 'A cat running')", label="Prompt")
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Kling AI Image-to-Video Generator",
    description="Upload a PNG/JPEG image to generate a 5-second video using Kling AI API."
)

# Launch the interface
if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)