import asyncio
import base64
import io
import logging
import os

import gradio as gr
import httpx
from PIL import Image

# Configure once at startup
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)

logger = logging.getLogger(__name__)


async def generate_async(images, prompt, variation, size):
    logger.info("giving it to BytePlus...")
    api_key = os.getenv("BYTEPLUS_API_KEY")
    base_url = os.getenv("BYTEPLUS_URL", "").rstrip("/")
    response_format = "b64_json"
    watermark = False

    images_input = []
    images_output = []

    # Convert uploaded images to base64
    for img in images or []:
        try:
            with open(img.name, "rb") as f:
                image_bytes = f.read()
                encoded = base64.b64encode(image_bytes).decode("utf-8")
                prefixed_b64 = f"data:image/png;base64,{encoded}"
                images_input.append(prefixed_b64)
        except Exception as e:
            logger.error(f"⚠️ Failed to process image {img.name}: {e}")

    try:
        model_name = "seedream-4-0-250828"
        request_data = {
            "model": model_name,
            "prompt": prompt,
            "response_format": response_format,
            "sequential_image_generation": "disabled",
            "size": size,
            "watermark": watermark,
        }

        if images_input:
            request_data["image"] = images_input

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        }
        logger.info("Sending request to BytePlus...")
        async with httpx.AsyncClient(timeout=120) as client:
            for _ in range(int(variation)):
                response = await client.post(
                    f"{base_url}/images/generations",
                    json=request_data,
                    headers=headers,
                )
                response.raise_for_status()
                result = response.json()

                # ✅ Fix: data is a list
                for item in result.get("data", []):
                    b64_str = item.get("b64_json")
                    if b64_str:
                        try:
                            image_data = base64.b64decode(b64_str)
                            image = Image.open(io.BytesIO(image_data))
                            images_output.append(image)
                        except Exception as e:
                            logger.warning(f"⚠️ Failed to decode base64 image: {e}")

        return images_output

    except Exception as e:
        logger.error(f"⚠️ Failed to process everything: {e}")
        return []


# Wrapper because Gradio doesn't await async functions
async def generate(images, prompt, variation, size):
    return await generate_async(images, prompt, variation, size)


# ------------------ Gradio UI ------------------ #
with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.Markdown("## 🔥 Multi-API Image-to-Image Generator")

    with gr.Row():
        # === Left Column ===
        with gr.Column(scale=1):
            # Upload input
            image_input = gr.File(
                label="Upload your reference images (optional)",
                file_count="multiple",
                file_types=["image"],
            )

            # Preview gallery for uploaded images
            input_preview = gr.Gallery(label="Preview", columns=3, height="auto")

            # Auto-preview when images uploaded
            image_input.change(
                lambda files: [f.name for f in files] if files else [],
                inputs=image_input,
                outputs=input_preview,
            )

            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Describe how you want to modify or generate the image...",
                lines=2,
            )

            variation_choice = gr.Dropdown(
                choices=[1, 2, 3, 4, 5],
                label="Number of Variations",
                value=1,
            )

            size_choice = gr.Dropdown(
                choices=[
                    "2048x2048",
                    "1728x2304",
                    "2304x1728",
                    "2560x1440",
                    "1440x2560",
                    "1664x2496",
                    "2496x1664",
                    "3024x1296",
                ],
                label="Output Size",
                value="2048x2048",
            )

            btn = gr.Button("🚀 Generate", variant="primary", size="lg")

        # === Right Column ===
        with gr.Column(scale=1):
            gallery = gr.Gallery(label="Generated Results", columns=2, height="auto")

    btn.click(
        fn=generate,
        inputs=[image_input, prompt_input, variation_choice, size_choice],
        outputs=gallery,
    )

demo.launch()