"""
Image Upscaler Comparison - Model Comparison
Powered by WaveSpeed AI - Auto-generated by Space Generator
"""

import gradio as gr
import requests
import time
from typing import Dict, Any

# ============ API Configuration ============
WAVESPEED_API_BASE = "https://api.wavespeed.ai/api/v3"
UPLOAD_ENDPOINT = "https://api.wavespeed.ai/api/v3/media/upload/binary"

MODEL_A_ENDPOINT = "wavespeed-ai/image-upscaler"
MODEL_A_NAME = "Image Upscaler"

MODEL_B_ENDPOINT = "wavespeed-ai/ultimate-image-upscaler"
MODEL_B_NAME = "Ultimate Upscaler"

POLL_INTERVAL = 1.5
POLL_MAX_SECONDS = 120





# ============ CSS ============
CUSTOM_CSS = """
/* ===== Base Styles ===== */
html, body, .gradio-container {
    background: linear-gradient(145deg, #f5f3ff 0%, #ede9fe 50%, #e0e7ff 100%) !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
}

.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 24px !important;
}

/* ===== Hero Section ===== */
.hero-container {
    text-align: center;
    padding: 40px 20px 30px;
}

.hero-badge {
    display: inline-block;
    background: linear-gradient(135deg, #06b6d4, #0891b2);
    padding: 10px 24px;
    border-radius: 50px;
    font-size: 0.7rem;
    color: #fff;
    font-weight: 700;
    letter-spacing: 1.5px;
    margin-bottom: 16px;
    box-shadow: 0 4px 20px rgba(139, 92, 246, 0.35);
}

.hero-title {
    font-size: 2.5rem;
    font-weight: 800;
    margin: 0 0 12px 0;
    background: linear-gradient(135deg, #6d28d9 0%, #06b6d4 50%, #a78bfa 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

.hero-desc {
    font-size: 1rem;
    color: #64748b;
    max-width: 100%;
    margin: 0 auto 16px;
    line-height: 1.6;
}

/* ===== Main Card ===== */
.main-card {
    background: #ffffff;
    border: 1px solid rgba(139, 92, 246, 0.1);
    border-radius: 20px;
    padding: 24px;
    margin-bottom: 16px;
    box-shadow: 0 4px 24px rgba(139, 92, 246, 0.08);
}

/* ===== Model Labels ===== */
.model-label {
    display: inline-block;
    padding: 6px 14px;
    border-radius: 8px;
    font-weight: 700;
    font-size: 0.85rem;
    margin-bottom: 12px;
}

.model-a-label {
    background: linear-gradient(135deg, #3b82f6, #2563eb);
    color: #fff;
}

.model-b-label {
    background: linear-gradient(135deg, #10b981, #059669);
    color: #fff;
}

/* ===== VS Divider ===== */
.vs-divider {
    display: flex;
    align-items: center;
    justify-content: center;
    padding: 8px 0;
}

.vs-badge {
    background: linear-gradient(135deg, #06b6d4, #0891b2);
    color: #fff;
    padding: 8px 16px;
    border-radius: 20px;
    font-weight: 800;
    font-size: 0.9rem;
    box-shadow: 0 4px 12px rgba(139, 92, 246, 0.3);
}

/* ===== API Key Section ===== */
.api-key-row {
    display: flex;
    align-items: center;
    justify-content: space-between;
    margin-bottom: 12px;
}

.api-key-label {
    display: flex;
    align-items: center;
    gap: 10px;
    color: #1e293b;
    font-weight: 700;
    font-size: 1rem;
}

.get-key-btn {
    padding: 10px 20px;
    background: linear-gradient(135deg, #06b6d4, #0891b2);
    border: none;
    border-radius: 10px;
    color: #fff !important;
    text-decoration: none;
    font-weight: 600;
    font-size: 0.85rem;
    transition: all 0.25s ease;
    box-shadow: 0 4px 12px rgba(139, 92, 246, 0.3);
}

.get-key-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 20px rgba(139, 92, 246, 0.4);
}

/* ===== Section Title ===== */
.section-title {
    color: #1e293b;
    font-weight: 700;
    font-size: 1rem;
    display: flex;
    align-items: center;
    gap: 10px;
    margin-bottom: 12px;
}

/* ===== Upload Area ===== */
.upload-area {
    border: 2px dashed rgba(139, 92, 246, 0.3) !important;
    border-radius: 16px !important;
    background: linear-gradient(145deg, #faf5ff 0%, #f5f3ff 100%) !important;
    min-height: 180px !important;
}

/* ===== Result Areas ===== */
.result-area-a {
    border: 2px solid rgba(59, 130, 246, 0.3) !important;
    border-radius: 16px !important;
    background: rgba(239, 246, 255, 0.5) !important;
    min-height: 200px !important;
}

.result-area-b {
    border: 2px solid rgba(16, 185, 129, 0.3) !important;
    border-radius: 16px !important;
    background: rgba(236, 253, 245, 0.5) !important;
    min-height: 200px !important;
}

/* ===== Button Styling ===== */
.compare-btn {
    width: 100%;
    margin-top: 16px !important;
    background: linear-gradient(135deg, #06b6d4, #0891b2) !important;
    border: none !important;
    color: #fff !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 16px 28px !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 16px rgba(139, 92, 246, 0.35) !important;
}

.compare-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 24px rgba(139, 92, 246, 0.45) !important;
}

/* ===== CTA Section ===== */
.cta-container {
    text-align: center;
    padding: 36px 28px;
    background: linear-gradient(135deg, #06b6d4 0%, #0891b2 50%, #6d28d9 100%);
    border-radius: 20px;
    margin-top: 8px;
    box-shadow: 0 8px 32px rgba(139, 92, 246, 0.35);
}

.cta-title {
    color: #fff;
    font-size: 1.4rem;
    font-weight: 800;
    margin: 0 0 8px 0;
}

.cta-desc {
    color: rgba(255, 255, 255, 0.9);
    font-size: 0.95rem;
    margin: 0 0 20px 0;
}

.cta-btn {
    display: inline-block;
    padding: 12px 32px;
    background: #fff;
    border-radius: 12px;
    color: #0891b2 !important;
    text-decoration: none;
    font-weight: 700;
    font-size: 0.95rem;
    box-shadow: 0 4px 16px rgba(0, 0, 0, 0.1);
}

.cta-btn:hover {
    transform: translateY(-3px);
    box-shadow: 0 8px 24px rgba(0, 0, 0, 0.15);
}

/* ===== Hide Elements ===== */
footer { display: none !important; }

/* ===== Input Styling ===== */
.gradio-container input[type="password"],
.gradio-container input[type="text"] {
    border: 2px solid #e2e8f0 !important;
    border-radius: 12px !important;
    padding: 12px 14px !important;
    font-size: 0.95rem !important;
}

.gradio-container input:focus {
    border-color: #06b6d4 !important;
    box-shadow: 0 0 0 3px rgba(139, 92, 246, 0.15) !important;
}
"""


# ============ API Functions ============
def upload_image(api_key: str, file_path: str) -> str:
    headers = {"Authorization": f"Bearer {api_key.strip()}"}
    with open(file_path, "rb") as f:
        resp = requests.post(UPLOAD_ENDPOINT, headers=headers, files={"file": f}, timeout=60)
    if resp.status_code == 401:
        raise Exception("Invalid API Key")
    elif resp.status_code >= 400:
        raise Exception(f"Upload failed: {resp.status_code}")
    data = resp.json()
    if data.get("code") != 200:
        raise Exception(data.get("message", "Upload failed"))
    return data.get("data", {}).get("download_url")


def call_api(api_key: str, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
    headers = {"Authorization": f"Bearer {api_key.strip()}", "Content-Type": "application/json"}
    resp = requests.post(f"{WAVESPEED_API_BASE}/{endpoint}", json=payload, headers=headers, timeout=30)
    if resp.status_code == 401:
        raise Exception("Invalid API Key")
    elif resp.status_code == 429:
        raise Exception("Quota exceeded")
    elif resp.status_code >= 400:
        raise Exception(f"API error: {resp.status_code}")
    data = resp.json()
    if data.get("code") != 200:
        raise Exception(data.get("message", "Unknown error"))
    return data.get("data", {})


def poll_result(api_key: str, request_id: str) -> Dict[str, Any]:
    headers = {"Authorization": f"Bearer {api_key.strip()}"}
    url = f"{WAVESPEED_API_BASE}/predictions/{request_id}/result"
    start_time = time.time()
    while time.time() - start_time < POLL_MAX_SECONDS:
        resp = requests.get(url, headers=headers, timeout=30)
        if resp.status_code >= 400:
            raise Exception("Failed to get result")
        result = resp.json().get("data", {})
        status = result.get("status", "")
        if status == "completed":
            return result
        elif status == "failed":
            raise Exception("Generation failed")
        time.sleep(POLL_INTERVAL)
    raise Exception("Timeout")


def run_model(api_key: str, endpoint: str, payload: Dict[str, Any]):
    """运行单个模型"""
    try:
        result = call_api(api_key, endpoint, payload)
        request_id = result.get("id")
        if not request_id:
            return None
        final_result = poll_result(api_key, request_id)
        outputs = final_result.get("outputs", [])
        if outputs:
            return outputs[0]
        return None
    except Exception as e:
        return None


def compare_models(api_key: str, image_path: str):
    """同时运行两个模型进行对比"""
    if not api_key or not api_key.strip():
        gr.Warning("Please enter your API Key")
        return None, None
    if not image_path:
        gr.Warning("Please upload an image")
        return None, None

    try:
        gr.Info("Uploading image...")
        image_url = upload_image(api_key, image_path)
    except Exception as e:
        gr.Warning(f"Upload failed: {e}")
        return None, None

    gr.Info(f"Running {MODEL_A_NAME}...")
    payload_a = {
            "image": image_url,
            "enable_base64_output": False,
            "enable_sync_mode": False
        }
    result_a = run_model(api_key, MODEL_A_ENDPOINT, payload_a)

    gr.Info(f"Running {MODEL_B_NAME}...")
    payload_b = {
            "image": image_url,
            "enable_base64_output": False,
            "enable_sync_mode": False
        }
    result_b = run_model(api_key, MODEL_B_ENDPOINT, payload_b)

    if result_a and result_b:
        gr.Info("Comparison complete!")
    elif result_a:
        gr.Warning(f"{MODEL_B_NAME} failed")
    elif result_b:
        gr.Warning(f"{MODEL_A_NAME} failed")
    else:
        gr.Warning("Both models failed")

    return result_a, result_b


# ============ Gradio UI ============
with gr.Blocks(css=CUSTOM_CSS, title="Image Upscaler Comparison") as demo:

    # Hero Section
    gr.HTML("""
        <div class="hero-container">
            <div class="hero-badge">WAVESPEED AI</div>
            <h1 class="hero-title">Image Upscaler Comparison</h1>
            <p class="hero-desc">Compare Image Upscaler vs Ultimate Upscaler side by side. Try on <a href="https://wavespeed.ai/models?utm_source=huggingface_space_upscaler_comparison" target="_blank" style="color: #8b5cf6; text-decoration: none; font-weight: 600;">wavespeed</a></p>
        </div>
    """)

    

    # API Key Card
    with gr.Group(elem_classes="main-card"):
        gr.HTML("""
            <div class="api-key-row">
                <span class="api-key-label">
                    <svg width="20" height="20" fill="#06b6d4" viewBox="0 0 20 20"><path fill-rule="evenodd" d="M18 8a6 6 0 01-7.743 5.743L10 14l-1 1-1 1H6v2H2v-4l4.257-4.257A6 6 0 1118 8zm-6-4a1 1 0 100 2 2 2 0 012 2 1 1 0 102 0 4 4 0 00-4-4z" clip-rule="evenodd"/></svg>
                    API Key
                </span>
                <a href="https://wavespeed.ai/models?utm_source=huggingface_space_upscaler_comparison" target="_blank" class="get-key-btn">Get API Key</a>
            </div>
        """)
        api_key_input = gr.Textbox(
            placeholder="Enter your WaveSpeed API key",
            type="password",
            show_label=False
        )

    # Input Section
    with gr.Group(elem_classes="main-card"):
        gr.HTML("""
            <div class="section-title">
                <svg width="20" height="20" fill="#06b6d4" viewBox="0 0 20 20"><path fill-rule="evenodd" d="M4 3a2 2 0 00-2 2v10a2 2 0 002 2h12a2 2 0 002-2V5a2 2 0 00-2-2H4zm12 12H4l4-8 3 6 2-4 3 6z" clip-rule="evenodd"/></svg>
                Input
            </div>
        """)
        image_input = gr.Image(
            label="Upload Image",
            type="filepath",
            source="upload",
            elem_classes="upload-area"
        )
        compare_btn = gr.Button("🔄 Compare Models", variant="primary", elem_classes="compare-btn")

    # Results Section
    with gr.Group(elem_classes="main-card"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML(f'<span class="model-label model-a-label">{MODEL_A_NAME}</span>')
                output_a = gr.Image(label="", interactive=False, elem_classes="result-area-a")

            with gr.Column(scale=1):
                gr.HTML(f'<span class="model-label model-b-label">{MODEL_B_NAME}</span>')
                output_b = gr.Image(label="", interactive=False, elem_classes="result-area-b")

    # CTA Section
    gr.HTML("""
        <div class="cta-container">
            <h3 class="cta-title">Want More Features?</h3>
            <p class="cta-desc">Higher resolutions, batch processing, and 700+ AI models</p>
            <a href="https://wavespeed.ai/models?utm_source=huggingface_space_upscaler_comparison" target="_blank" class="cta-btn">
                Explore WaveSpeed.ai
            </a>
        </div>
    """)

    # Event binding
    compare_btn.click(
        fn=compare_models,
        inputs=[api_key_input, image_input],
        outputs=[output_a, output_b],
    )


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