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

# ================== CONFIG ==================
ACCESS_KEY_ID = "YOUR_ACCESS_KEY_ID"
ACCESS_KEY_SECRET = "YOUR_ACCESS_KEY_SECRET"
API_URL = "https://api-singapore.klingai.com/v1/images/generations"

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


# ================== AUTH ==================
def generate_jwt_token():
    headers = {"alg": "HS256", "typ": "JWT"}
    payload = {
        "iss": ACCESS_KEY_ID,
        "exp": int(time.time()) + 1800,
        "nbf": int(time.time()) - 5
    }
    return jwt.encode(payload, ACCESS_KEY_SECRET, headers=headers)


# ================== 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 reading image: {e}"

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

    payload = {
        "model_name": "kling-v2-1",           # ✅ per docs
        "prompt": prompt or "Transform face into artistic style",
        "image": image_base64,                # ✅ base64 only, no prefix
        "aspect_ratio": "1:1",
        "n": 1,                                # number of images
        "resolution": "1k"
    }

    try:
        logger.debug(f"POST {API_URL} payload={payload.keys()}")
        response = requests.post(API_URL, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        data = response.json()

        if data.get("code") != 0:
            return None, f"API Error: {data.get('message')}"

        task_id = data["data"]["task_id"]
        logger.debug(f"Task ID: {task_id}")

        # Poll task status
        status_url = f"{API_URL}/{task_id}"
        for _ in range(60):  # up to ~5 minutes
            time.sleep(5)
            status_resp = requests.get(status_url, headers=headers, timeout=30)
            status_data = status_resp.json()

            if status_data.get("data", {}).get("task_status") == "succeed":
                images = status_data["data"]["task_result"]["images"]
                if images:
                    img_url = images[0]["url"]
                    img_resp = requests.get(img_url)
                    img_resp.raise_for_status()
                    out_path = Path("/tmp/output.png")
                    with open(out_path, "wb") as f:
                        f.write(img_resp.content)
                    return str(out_path), None
                else:
                    return None, "No image returned."

            elif status_data.get("data", {}).get("task_status") == "failed":
                return None, f"Generation failed: {status_data['data'].get('task_status_msg')}"

        return None, "Timed out waiting for result."

    except requests.exceptions.HTTPError as e:
        return None, f"HTTP Error: {e}"
    except Exception as e:
        return None, f"Unexpected Error: {e}"


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


with gr.Blocks() as iface:
    gr.Markdown("# Kling AI Image Generator — Official API Format")

    with gr.Row():
        with gr.Column():
            img_in = gr.Image(type="filepath", label="Upload Image")
            prompt_in = gr.Textbox(lines=2, placeholder="Enter prompt", label="Prompt")
            btn = gr.Button("Generate")
        with gr.Column():
            img_out = gr.Image(label="Generated Image")
            file_out = gr.File(label="Download Image")
            err_out = gr.Textbox(label="Error / Status", interactive=False)

    btn.click(fn=chatbot_interface, inputs=[img_in, prompt_in], outputs=[img_out, file_out, err_out])

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