import gradio as gr
import spaces  # Necessário para o decorador @spaces.GPU (caso esteja usando Hugging Face Spaces)
import os
import torch
from datetime import datetime
from PIL import Image
import boto3
from botocore.exceptions import NoCredentialsError
from dotenv import load_dotenv
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler

# Carregar variáveis de ambiente do arquivo .env
load_dotenv()

# Configurações do AWS S3
AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')
AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
AWS_BUCKET_NAME = os.getenv('AWS_BUCKET_NAME')
AWS_REGION = os.getenv('AWS_REGION')
HF_TOKEN = os.getenv('HF_TOKEN')  # Token da Hugging Face

# Inicializar cliente S3
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY,
    aws_secret_access_key=AWS_SECRET_KEY,
    region_name=AWS_REGION
)

# Configuração do pipeline para "character"
character_pipe = DiffusionPipeline.from_pretrained(
    "cagliostrolab/animagine-xl-3.1",
    torch_dtype=torch.float16,
    use_safetensors=True,
    use_auth_token=HF_TOKEN  # Inclui o token aqui
)
character_pipe.scheduler = EulerDiscreteScheduler.from_config(character_pipe.scheduler.config)

# Configuração do pipeline para "item"
item_pipe = DiffusionPipeline.from_pretrained(
    "openart-custom/DynaVisionXL",
    torch_dtype=torch.float16,
    use_safetensors=True,
    use_auth_token=HF_TOKEN  # Inclui o token aqui
)
item_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(item_pipe.scheduler.config)

# Função de geração de imagem com alocação de GPU (através do decorador do Hugging Face Spaces)
@spaces.GPU(duration=60)  # Aloca a GPU somente durante a execução desta função
def generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
    if model_type == "character":
        pipe = character_pipe
        default_prompt = "1girl, souji okita, fate series, solo, upper body, bedroom, night, seducing, (sexy clothes)"
        default_negative_prompt = ("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, "
                                   "low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, "
                                   "signature, extra digits, artistic error, username, scan, [abstract]")
    elif model_type == "item":
        pipe = item_pipe
        default_prompt = "great sword, runes on blade, acid on blade, weapon, (((item)))"
        default_negative_prompt = "1girl, girl, man, boy, 1man, men, girls"
    else:
        return "Tipo inválido. Escolha entre 'character' ou 'item'."

    # Se o usuário fornecer prompt, utiliza-o; caso contrário, usa o padrão
    final_prompt = prompt if prompt else default_prompt
    final_negative_prompt = negative_prompt if negative_prompt else default_negative_prompt

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

    # Geração da imagem
    result = pipe(
        prompt=final_prompt,
        negative_prompt=final_negative_prompt,
        width=int(width),
        height=int(height),
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(num_inference_steps)
    )
    image = result.images[0]

    # Salva a imagem em um arquivo temporário
    temp_file = "/tmp/generated_image.png"
    image.save(temp_file)

    # Faz upload para o AWS S3
    file_name = datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
    try:
        s3_client.upload_file(temp_file, AWS_BUCKET_NAME, file_name)
        s3_url = f"https://{AWS_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{file_name}"
        return s3_url
    except NoCredentialsError:
        return "Credenciais não disponíveis"

# Função que integra a geração via Gradio
def gradio_generate(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps):
    return generate_image(model_type, prompt, negative_prompt, width, height, guidance_scale, num_inference_steps)

# Definindo os componentes de entrada utilizando a API atual do Gradio
model_type_input = gr.Dropdown(choices=["character", "item"], value="character", label="Model Type")
prompt_input = gr.Textbox(lines=2, placeholder="Digite o prompt (deixe vazio para o padrão)", label="Prompt")
negative_prompt_input = gr.Textbox(lines=2, placeholder="Digite o negative prompt (deixe vazio para o padrão)", label="Negative Prompt")
width_input = gr.Number(value=832, label="Width")
height_input = gr.Number(value=1216, label="Height")
guidance_scale_input = gr.Number(value=10.0, label="Guidance Scale")
num_inference_steps_input = gr.Number(value=100, label="Number of Inference Steps")

# Criação da interface Gradio
iface = gr.Interface(
    fn=gradio_generate,
    inputs=[
        model_type_input,
        prompt_input,
        negative_prompt_input,
        width_input,
        height_input,
        guidance_scale_input,
        num_inference_steps_input,
    ],
    outputs="text",
    title="Image Generation API",
    description="Gere imagens usando modelos de difusão e faça upload para o AWS S3."
)

if __name__ == "__main__":
    iface.launch()