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

# Configuration - REPLACE WITH YOUR ACTUAL CREDENTIALS
ACCESS_KEY_ID = "AFyHfnQATghFdCMyAG3gRPbNY4TNKFGB"
ACCESS_KEY_SECRET = "TTepeLyBterLNM3brYPGmdndBnnyKJBA"
API_BASE_URL = "https://api-singapore.klingai.com"
ENDPOINT = f"{API_BASE_URL}/v1/images/generations"  # Image-to-image endpoint

def generate_jwt_token():
    """Generate authentication token"""
    payload = {
        "iss": ACCESS_KEY_ID,
        "exp": int(time.time()) + 1800,  # 30 min expiration
        "nbf": int(time.time()) - 5      # Not before 5 sec ago
    }
    return jwt.encode(payload, ACCESS_KEY_SECRET, algorithm="HS256")

def process_image(image_path, prompt):
    """Core image processing function"""
    try:
        # 1. Validate image
        if not os.path.exists(image_path):
            return None, "Image file not found"
        
        if os.path.getsize(image_path) > 10 * 1024 * 1024:  # 10MB
            return None, "Image too large (max 10MB)"
        
        # 2. Prepare image
        with open(image_path, "rb") as f:
            image_base64 = base64.b64encode(f.read()).decode('utf-8')
        
        # 3. API Request
        headers = {
            "Authorization": f"Bearer {generate_jwt_token()}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model_name": "kling-v2.1",
            "prompt": prompt,
            "image": image_base64,
            "image_reference": "face",
            "image_fidelity": 0.97,
            "human_fidelity": 0.97,
            "aspect_ratio": "1:1",
            "n": 1
        }
        
        response = requests.post(ENDPOINT, json=payload, headers=headers)
        
        # 4. Handle response
        if response.status_code != 200:
            return None, f"API Error: {response.text}"
        
        data = response.json()
        if data.get("code") != 0:
            return None, f"API Error: {data.get('message', 'Unknown error')}"
        
        task_id = data["data"]["task_id"]
        
        # 5. Check task status (max 3 minutes)
        for _ in range(18):  # 18 attempts × 10 seconds = 3 minutes
            time.sleep(10)
            status_response = requests.get(
                f"{API_BASE_URL}/v1/images/generations/{task_id}",
                headers=headers
            )
            status_data = status_response.json()
            
            if status_data["data"]["task_status"] == "succeed":
                image_url = status_data["data"]["task_result"]["images"][0]["url"]
                img_data = requests.get(image_url).content
                output_path = f"/tmp/result_{task_id}.png"
                with open(output_path, "wb") as f:
                    f.write(img_data)
                return output_path, None
                
            elif status_data["data"]["task_status"] in ("failed", "canceled"):
                return None, status_data["data"].get("task_status_msg", "Task failed")
        
        return None, "Processing timed out"
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# 🖼️ Face Style Transformer")
    gr.Markdown("Upload a clear face photo and describe your desired style")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Face Photo")
            prompt_input = gr.Textbox(label="Style Prompt", 
                                   placeholder="e.g. 'anime character', 'oil painting'")
            generate_btn = gr.Button("Transform", variant="primary")
            
            gr.Markdown("### Requirements:")
            gr.Markdown("""
            - Clear frontal face photo
            - Single person only
            - Max 10MB (JPG/PNG)
            - Min 300x300 resolution
            """)
            
        with gr.Column():
            output_image = gr.Image(label="Result", interactive=False)
            output_file = gr.File(label="Download Result")
            status_output = gr.Textbox(label="Status")
    
    generate_btn.click(
        fn=lambda img, prompt: process_image(img, prompt) + (None,),
        inputs=[image_input, prompt_input],
        outputs=[output_image, output_file, status_output]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)