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-image"

def generate_jwt_token():
    """Generate JWT token for Kling AI API authentication."""
    headers = {
        "alg": "HS256",
        "typ": "JWT"
    }
    payload = {
        "iss": ACCESS_KEY_ID,
        "exp": int(time.time()) + 1800,  # Token expires in 30 minutes
        "nbf": int(time.time()) - 5  # Effective 5 seconds before current time
    }
    token = jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)
    return token

def generate_image(image, prompt=""):
    """
    Call Kling AI API to generate an image from an input image.
    
    Args:
        image: Uploaded image file (from Gradio)
        prompt (str): Optional text prompt to guide image transformation
    
    Returns:
        str: Path to the generated image or error message
    """
    if not image:
        return "Error: Please upload a valid image (PNG/JPEG, <10 MB, ≥512x512 pixels)."
    
    # 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. Ensure it’s a valid PNG/JPEG. Details: {str(e)}"
    
    headers = {
        "Authorization": f"Bearer {generate_jwt_token()}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "image": image_base64,
        "prompt": prompt or "Transform the image into a cartoon style",
        "strength": 0.7,  # Controls transformation intensity (0.0 to 1.0)
        "output_format": "png"
    }
    
    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") or data.get("id")
        if not task_id:
            return "Error: No task ID returned. Check API endpoint or keys at https://klingai.com/global/dev."
        
        # 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()
            status = status_data.get("status")
            if status == "succeeded":
                image_url = status_data.get("image_url") or status_data.get("result", {}).get("image_url") or status_data.get("output")
                if not image_url:
                    return "Error: No image URL in API response. Check API documentation."
                # Download the image
                image_response = requests.get(image_url, timeout=30)
                image_response.raise_for_status()
                output_path = Path("output_image.png")
                with open(output_path, "wb") as f:
                    f.write(image_response.content)
                return str(output_path)
            elif status == "failed":
                return "Error: Image generation failed. Avoid NSFW, celebrities, or sensitive content (e.g., politics, violence)."
            elif status == "processing" and status_data.get("progress", 0) >= 0.99:
                return "Error: Generation stuck at 99%. Check account credits or upgrade to a paid plan at https://klingai.com."
            time.sleep(5)
        
        return "Error: Image generation timed out. Try during off-peak hours or check account status."
    
    except requests.exceptions.HTTPError as e:
        status_code = e.response.status_code if e.response else None
        error_data = e.response.json() if e.response and e.response.text else {}
        service_code = error_data.get("code", 0)
        
        if status_code == 401:
            if service_code == 1000:
                return "Error: Authentication failed (1000). Verify API keys are correct."
            elif service_code == 1001:
                return "Error: Authorization empty (1001). Ensure JWT token is included."
            elif service_code == 1002:
                return "Error: Authorization invalid (1002). Check key format."
            elif service_code == 1003:
                return "Error: Authorization not yet valid (1003). Adjust token start time (nbf)."
            elif service_code == 1004:
                return "Error: Authorization expired (1004). Generate a new token."
        elif status_code == 429:
            if service_code == 1100:
                return "Error: Account exception (1100). Verify account configuration at https://klingai.com."
            elif service_code == 1101:
                return "Error: Account in arrears (1101). Recharge your account."
            elif service_code == 1102:
                return "Error: Resource pack depleted or expired (1102). Purchase additional resources."
            elif service_code in (1302, 1303):
                return "Error: Rate limit exceeded (1302/1303). Reduce request frequency or contact support."
            elif service_code == 1304:
                return "Error: IP whitelisting issue (1304). Contact Kling AI support."
        elif status_code == 403 and service_code == 1103:
            return "Error: Unauthorized access (1103). Verify account permissions."
        elif status_code == 400:
            if service_code == 1200:
                return "Error: Invalid request parameters (1200). Check payload format."
            elif service_code == 1201:
                return "Error: Invalid parameters (1201). Use PNG/JPEG (<10 MB, ≥512x512) and a simple prompt."
            elif service_code == 1300:
                return "Error: Platform policy triggered (1300). Check input content."
            elif service_code == 1301:
                return "Error: Content security policy triggered (1301). Avoid NSFW or sensitive content."
        elif status_code == 404:
            if service_code == 1202:
                return "Error: Invalid request method (1202). Check https://klingai.com/global/dev for correct endpoint."
            elif service_code == 1203:
                return "Error: Resource does not exist (1203). Verify model or endpoint."
        elif status_code in (500, 503, 504):
            if service_code == 5000:
                return "Error: Server internal error (5000). Try again later."
            elif service_code == 5001:
                return "Error: Server unavailable due to maintenance (5001). Try again later."
            elif service_code == 5002:
                return "Error: Server timeout (5002). Try during off-peak hours."
        return f"Error: API request failed. HTTP {status_code}, Service Code {service_code}. Details: {str(e)}"
    except requests.exceptions.RequestException as e:
        return f"Error: Network issue. Ensure stable internet, disable VPN, and try again. Details: {str(e)}"

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

# Define Gradio interface
iface = gr.Interface(
    fn=chatbot_interface,
    inputs=[
        gr.Image(type="filepath", label="Upload Image (PNG/JPEG, <10 MB, ≥512x512)"),
        gr.Textbox(lines=2, placeholder="Enter an optional prompt (e.g., 'Turn this into a cartoon')", label="Prompt")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Kling AI Image-to-Image Generator",
    description="Upload a PNG/JPEG image (<10 MB, ≥512x512 pixels) to generate a transformed image using Kling AI API. Avoid NSFW, celebrities, or sensitive content."
)

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