import gradio as gr
import base64
import mimetypes
import os
import google-genai as genai
from google-genai import types
from PIL import Image
import io


# ==========================
# CONFIG GOOGLE API KEY
# ==========================
API_KEY = os.environ.get("GEMINI_API_KEY")
client = genai.Client(api_key=API_KEY)


# ==========================
# FUNÇÃO DE GERAÇÃO (COM A NOVA SDK)
# ==========================
def generate_image(prompt, negative_prompt, resolution):

    if not API_KEY:
        return None, "❌ API Key GEMINI_API_KEY não configurada no HuggingFace."

    try:
        # Construir prompt completo
        full_prompt = prompt
        if negative_prompt:
            full_prompt += f"\n\nEvitar: {negative_prompt}"

        contents = [
            types.Content(
                role="user",
                parts=[types.Part.from_text(text=full_prompt)],
            )
        ]

        # Resolução suportada pelo modelo
        image_res = {"1K": "1K", "2K": "2K", "4K": "4K"}.get(resolution, "1K")

        config = types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
            image_config=types.ImageConfig(image_size=image_res),
            tools=[types.Tool(googleSearch=types.GoogleSearch())],
        )

        # STREAM da geração (oficial)
        for chunk in client.models.generate_content_stream(
            model="gemini-3-pro-image-preview",
            contents=contents,
            config=config,
        ):
            if (
                chunk.candidates
                and chunk.candidates[0].content
                and chunk.candidates[0].content.parts
            ):

                part = chunk.candidates[0].content.parts[0]

                # ==========================
                # 1. inline_data → imagem real
                # ==========================
                if hasattr(part, "inline_data") and part.inline_data:
                    mime = part.inline_data.mime_type
                    if mime and mime.startswith("image"):
                        data = part.inline_data.data
                        img = Image.open(io.BytesIO(data))
                        return img, "✅ Imagem gerada com sucesso!"

                # ==========================
                # 2. fallback: part.image
                # ==========================
                if hasattr(part, "image") and part.image:
                    try:
                        img = Image.open(io.BytesIO(part.image))
                        return img, "✅ Imagem gerada (fallback image)."
                    except:
                        pass

                # ==========================
                # 3. fallback: part.blob
                # ==========================
                if hasattr(part, "blob") and part.blob:
                    try:
                        img = Image.open(io.BytesIO(part.blob))
                        return img, "✅ Imagem gerada (fallback blob)."
                    except:
                        pass

        return None, "❌ O modelo respondeu, mas não retornou imagem."

    except Exception as e:
        return None, f"❌ Erro: {str(e)}"


# ==========================
# EXEMPLOS DE PROMPT
# ==========================
examples = [
    [
        "Cinematic portrait of a woman with red hair, soft light, 85mm, ultra realistic, 8k",
        "blurry, distorted, ugly, low quality",
        "1K",
    ],
    [
        "Cyberpunk futuristic city, neon rain, flying cars, ultrarealistic, night mood",
        "daylight, cartoon, lowres",
        "2K",
    ],
    [
        "Mystical forest, god rays, fog, moss rocks, photorealistic nature",
        "urban, artificial",
        "1K",
    ],
]


# ==========================
# INTERFACE GRADIO
# ==========================
with gr.Blocks() as demo:

    # CSS ☑ estética Leicam
    gr.HTML("""
    <style>
        .gradio-container {
            font-family: 'Inter','Manrope',sans-serif;
        }
        .title {
            text-align: center;
            font-size: 2.6em;
            font-weight: 800;
            margin-bottom: 0.3em;
            background: linear-gradient(135deg,#39FF14 0%,#00CC11 100%);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
        }
        .subtitle {
            text-align: center;
            font-size: 1.1em;
            color:#6b7280;
            margin-bottom: 2em;
        }
    </style>
    """)

    # HEADER
    gr.HTML("""
        <div style='text-align:center;margin-bottom:30px;'>
            <h1 class='title'>🎨 Gerador Ultra-Realista (Nano Banana Pro)</h1>
            <p class='subtitle'>Gemini 3 Pro Image Preview — Imagens de nível profissional</p>
        </div>
    """)

    # LAYOUT
    with gr.Row():

        # COLUNA ESQUERDA
        with gr.Column():

            prompt = gr.Textbox(
                label="📝 Prompt",
                lines=5,
                placeholder="Descreva a imagem desejada com detalhes..."
            )

            negative_prompt = gr.Textbox(
                label="🚫 Negative Prompt",
                value="blurry, distorted, ugly, deformed",
                lines=3
            )

            resolution = gr.Dropdown(
                label="📐 Resolução",
                choices=["1K", "2K", "4K"],
                value="1K"
            )

            btn = gr.Button("✨ Gerar imagem", variant="primary")

        # COLUNA DIREITA
        with gr.Column():

            output_image = gr.Image(
                type="pil",
                height=600,
                label="Imagem Gerada"
            )

            output_text = gr.Textbox(
                label="Status",
                lines=5,
                interactive=False
            )

    # EXEMPLOS
    gr.Markdown("### 📚 Exemplos")
    gr.Examples(
        examples=examples,
        inputs=[prompt, negative_prompt, resolution]
    )

    # FOOTER
    gr.HTML("""
        <div style='text-align:center;margin-top:50px;padding:20px;border-top:1px solid #e5e7eb;'>
            <strong>Leicam · Tech</strong><br>
            <span style='color:#9ca3af;font-size:12px;'>© 2025 Todos os direitos reservados.</span>
        </div>
    """)

    # EVENTO
    btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, resolution],
        outputs=[output_image, output_text],
    )


# RODAR LOCALMENTE
if __name__ == "__main__":
    demo.launch()