import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random
import os
import json
import io
import uuid
from gradio_client import Client as client_gradio
from supabase import create_client, Client
from datetime import datetime

# Inicializa Supabase
url: str = os.getenv('SUPABASE_URL')
key: str = os.getenv('SUPABASE_KEY')
supabase: Client = create_client(url, key)

# Obtém token da Hugging Face
hf_token = os.getenv("HF_TOKEN")

# Inicializa o modelo base FLUX.1-dev
base_model = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16, use_safetensors=True)

# Move o modelo para GPU
pipe.to("cuda")

# Definição dos LoRA e Trigger Words
lora_models = {
    "AndroFlux": {
        "repo": "vcollos/Nanda",
        "weights": "lora.safetensors",
        "trigger_word": ""  # Sem trigger word específica
    },
    "vgnCollos": {
        "repo": "vcollos/VitorCollos",
        "weights": "Vitor.safetensors",
        "trigger_word": "A photo of Vitor,"
    }
}

# Carrega os LoRAs
for name, details in lora_models.items():
    try:
        pipe.load_lora_weights(details["repo"], weight_name=details["weights"], adapter_name=name)
        print(f"✅ LoRA {name} carregado")
    except Exception as e:
        print(f"❌ Erro ao carregar o LoRA {name}: {e}")

# Define seed máximo
MAX_SEED = 2**32 - 1

def upload_image_to_supabase(image, filename):
    """ Faz upload da imagem para o Supabase Storage e retorna a URL pública. """
    img_bytes = io.BytesIO()
    image.save(img_bytes, format="PNG")
    img_bytes.seek(0)  # Move para o início do arquivo

    storage_path = f"images/{filename}"

    try:
        # Upload da imagem
        supabase.storage.from_("images").upload(storage_path, img_bytes.getvalue(), {"content-type": "image/png"})

        # Retorna a URL pública
        base_url = f"{url}/storage/v1/object/public/images"
        return f"{base_url}/{filename}"
    except Exception as e:
        print(f"❌ Erro no upload da imagem: {e}")
        return None

@spaces.GPU(duration=80)
def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale_1, lora_scale_2, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # Aplica os dois LoRAs combinados
    pipe.set_adapters(["AndroFlux", "vgnCollos"], adapter_weights=[lora_scale_1, lora_scale_2])

    # Adiciona trigger words apenas se vgnCollos estiver ativado
    if lora_scale_2 > 0:
        prompt = f"{lora_models['vgnCollos']['trigger_word']} {prompt}"

    # Gera a imagem
    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    # Define um nome único para a imagem
    filename = f"image_{seed}_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.png"

    try:
        image_url = upload_image_to_supabase(image, filename)
        print(f"✅ Imagem salva no Supabase: {image_url}")
    except Exception as e:
        print(f"❌ Erro ao fazer upload da imagem: {e}")
        image_url = None

    # Salva os metadados no banco de dados Supabase
    try:
        supabase.table("images").insert({
            "id": str(uuid.uuid4()),  # ID único
            "prompt": prompt,
            "cfg_scale": cfg_scale,
            "steps": steps,
            "seed": seed,
            "lora_scale_1": lora_scale_1,
            "lora_scale_2": lora_scale_2,
            "image_url": image_url,
            "created_at": datetime.utcnow().isoformat()
        }).execute()
        print("✅ Metadados salvos no Supabase")
    except Exception as e:
        print(f"❌ Erro ao salvar metadados no Supabase: {e}")

    return image, seed

# Interface Gradio
gr_theme = os.getenv("THEME")
with gr.Blocks(theme=gr_theme) as app:
    gr.Markdown("# Androflux Image Generator")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.TextArea(label="Prompt", placeholder="Digite um prompt (máx 77 caracteres)", lines=3)
            generate_button = gr.Button("Gerar")
            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
            steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=32)
            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=768)  
            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=1024)  
            randomize_seed = gr.Checkbox(False, label="Randomize seed")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=556215326)

            # Sliders para os pesos dos LoRAs
            lora_scale_1 = gr.Slider(label="LoRA Scale (AndroFlux)", minimum=0, maximum=1, step=0.01, value=0.1)
            lora_scale_2 = gr.Slider(label="LoRA Scale (vgnCollos)", minimum=0, maximum=1, step=0.01, value=1)

        with gr.Column(scale=2):
            result = gr.Image(label="Generated Image")
            gr.Markdown("Gere imagens usando Collos LoRA e um prompt de texto.")

    # Botão para gerar imagem combinando os LoRAs
    generate_button.click(
        run_lora,
        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale_1, lora_scale_2],
        outputs=[result, seed],
    )

app.queue()
app.launch(share=True)