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

# ================== CONFIGURATION ==================
ACCESS_KEY_ID = "AGBGmadNd9hakFYfahytyQQJtN8CJmDJ"
ACCESS_KEY_SECRET = "dp3pAe4PpdmnAHCAPgEd3PyLmBQrkMde"
API_URL = "https://api-singapore.klingai.com/v1/images/generations"  # New unified endpoint

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


# ================== AUTH ==================
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,
        "nbf": int(time.time()) - 5
    }
    token = jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)
    return token


# ================== IMAGE GENERATION ==================
def generate_image(image, prompt=""):
    if not image:
        return None, "Error: Please upload a valid face image."

    try:
        with open(image, "rb") as img_file:
            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
    except Exception as e:
        return None, f"Error: Failed to process image. Details: {str(e)}"

    headers = {
        "Authorization": f"Bearer {generate_jwt_token()}",
        "Content-Type": "application/json"
    }

    payload = {
        "model_name": "kolors-v2.1",
        "image_base64": [image_base64],  # New format: array of base64 images
        "prompt": prompt or "Transform the face into a cartoon style while preserving identity",
        "strength": 0.97,
        "output_format": "png",
        "n": 1,
        "aspect_ratio": "1:1"
    }

    try:
        logger.debug(f"Sending POST request to {API_URL}")
        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("data", {}).get("task_id")
        if not task_id:
            return None, "Error: No task ID returned from API."

        # ================== POLLING ==================
        status_url = f"https://api-singapore.klingai.com/v1/tasks/{task_id}"
        for _ in range(60):
            time.sleep(5)
            status_response = requests.get(status_url, headers=headers, timeout=30)
            if status_response.status_code == 404:
                continue
            status_data = status_response.json()
            status = status_data.get("task_status") or status_data.get("status")

            if status in ("succeeded", "completed"):
                image_url = (
                    status_data.get("image_url")
                    or status_data.get("data", {}).get("image")
                    or status_data.get("data", {}).get("images", [{}])[0].get("url")
                )
                if not image_url:
                    return None, "Error: No image URL in API response."

                image_response = requests.get(image_url, timeout=30)
                image_response.raise_for_status()
                output_path = Path("/tmp/output_image.png")
                with open(output_path, "wb") as f:
                    f.write(image_response.content)
                return str(output_path), None

            elif status == "failed":
                return None, "Error: Image generation failed."

        return None, "Error: Image generation timed out."

    except requests.exceptions.HTTPError as e:
        return None, f"HTTP Error: {str(e)}"
    except requests.exceptions.RequestException as e:
        return None, f"Network Error: {str(e)}"


# ================== GRADIO INTERFACE ==================
def chatbot_interface(image, prompt):
    output_path, error = generate_image(image, prompt)
    if error:
        return None, None, error
    return output_path, output_path, None


with gr.Blocks() as iface:
    gr.Markdown("# Kling AI Image-to-Image Generator (Kolors v2.1)")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Face Image")
            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt", label="Prompt")
            generate_button = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
            output_file = gr.File(label="Download Image")
            error_message = gr.Textbox(label="Status/Error", interactive=False)

    generate_button.click(
        fn=chatbot_interface,
        inputs=[image_input, prompt_input],
        outputs=[output_image, output_file, error_message]
    )

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