import os
import sys
import uuid
import json
from typing import List, Dict, Any, Optional

import torch
import gradio as gr
from PIL import Image
from omegaconf import OmegaConf, DictConfig

# --- 1. CONFIGURAÇÃO E IMPORTS ---

# Adiciona o diretório do VINCIE ao Python Path para importação de módulos.
VINCIE_DIR = os.getenv("VINCIE_DIR", "/app/VINCIE")
if VINCIE_DIR not in sys.path:
    sys.path.append(VINCIE_DIR)

try:
    from generate import VINCIEGenerator
    from common.config import load_config
    from common.seed import shift_seed
except ImportError:
    print(f"FATAL: Não foi possível importar os módulos do VINCIE. "
          f"Verifique se o repositório está em '{VINCIE_DIR}'.")
    raise

# --- 2. INICIALIZAÇÃO DO MODELO (SINGLETON) ---

MODEL: Optional[VINCIEGenerator] = None
DEVICE: Optional[torch.device] = None

def setup_model():
    """
    Inicializa e configura o modelo VINCIE em uma única GPU.
    Esta função é chamada uma vez no início da aplicação.
    """
    global MODEL, DEVICE

    if not torch.cuda.is_available():
        raise RuntimeError("FATAL: Nenhuma GPU compatível com CUDA foi encontrada.")

    num_gpus = torch.cuda.device_count()
    if num_gpus == 0:
        raise RuntimeError("FATAL: Nenhuma GPU foi detectada pelo PyTorch.")
    
    print(f"INFO: Detectadas {num_gpus} GPUs. A aplicação usará 'cuda:0'.")
    DEVICE = torch.device("cuda:0")
    torch.cuda.set_device(DEVICE)

    config_path = os.path.join(VINCIE_DIR, "configs/generate.yaml")
    print(f"INFO: Carregando e resolvendo configuração de '{config_path}'...")
    config = load_config(config_path, [])
    
    print("INFO: Instanciando VINCIEGenerator...")
    model_instance = VINCIEGenerator(config)
    
    print("INFO: Executando sequência de inicialização do VINCIE...")
    model_instance.configure_persistence()
    model_instance.configure_models()
    model_instance.configure_diffusion()
    
    if not hasattr(model_instance, 'dit'):
        raise RuntimeError("FATAL: Falha ao inicializar o componente DiT do modelo.")

    # Move todos os componentes para o dispositivo principal.
    model_instance.dit.to(DEVICE)
    model_instance.vae.to(DEVICE)
    model_instance.text_encoder.to(DEVICE)
    
    MODEL = model_instance
    print(f"✅ SUCESSO: Modelo VINCIE pronto para uso na GPU {DEVICE}.")

# --- 3. LÓGICAS DE INFERÊNCIA ---

def _execute_vincie_logic(
    prompt_config: DictConfig,
    steps: int,
    cfg_scale: float,
    seed: int,
    pad_img_placeholder: bool,
    resolution: int
) -> Image.Image:
    """
    Função central que executa a pipeline de inferência do VINCIE.
    """
    # Salva o estado original da configuração para restaurá-lo depois.
    original_config_state = {
        "steps": MODEL.config.diffusion.timesteps.sampling.steps,
        "seed": MODEL.config.generation.seed,
        "pad": MODEL.config.generation.pad_img_placehoder,
        "resolution": MODEL.config.generation.resolution,
    }
    
    try:
        OmegaConf.set_readonly(MODEL.config, False)
        
        # 1. Aplica configurações dinâmicas da UI.
        MODEL.config.diffusion.timesteps.sampling.steps = int(steps)
        MODEL.configure_diffusion()  # Recria o sampler com os novos passos.
        
        current_seed = seed if seed != -1 else torch.randint(0, 2**32 - 1, (1,)).item()
        MODEL.config.generation.seed = shift_seed(current_seed, 0)
        MODEL.config.generation.pad_img_placehoder = pad_img_placeholder
        MODEL.config.generation.resolution = int(resolution)

        # Log detalhado dos argumentos que serão enviados para a pipeline.
        _log_pipeline_args(prompt_config, steps, cfg_scale, MODEL.config.generation.seed, resolution, pad_img_placeholder)

        # 2. Prepara as entradas para o modelo.
        text_pos, condition, noise, _, _ = MODEL.prepare_input(
            prompt=prompt_config, repeat_idx=0, device=DEVICE
        )
        
        # 3. Executa a inferência.
        with torch.no_grad():
            samples = MODEL.inference(
                noises=[noise],
                conditions=[condition],
                texts_pos=[text_pos],
                texts_neg=[MODEL.config.generation.negative_prompt],
                cfg_scale=cfg_scale
            )
            
        if not samples:
            raise RuntimeError("A inferência do modelo não produziu resultados.")
            
        # 4. Processa a saída para formato de imagem.
        output_tensor = samples[0][:, -1, :, :]
        output_image_np = output_tensor.clip(-1, 1).add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
        return Image.fromarray(output_image_np)

    finally:
        # 5. Restaura a configuração original para garantir consistência entre chamadas.
        OmegaConf.set_readonly(MODEL.config, False)
        for key, value in original_config_state.items():
            if key == "steps":
                MODEL.config.diffusion.timesteps.sampling.steps = value
            else:
                OmegaConf.update(MODEL.config.generation, key, value, merge=True)
        OmegaConf.set_readonly(MODEL.config, True)
        MODEL.configure_diffusion() # Restaura o sampler padrão.

def run_single_turn_inference(
    input_image: str, prompt: str, aspect_ratio: str, resolution: int, steps: int, cfg_scale: float, seed: int
) -> Image.Image:
    """Handler para a aba 'Edição Simples'."""
    if not all([input_image, prompt]):
        raise gr.Error("É necessário fornecer uma imagem de entrada e um prompt.")
    
    prompt_config = OmegaConf.create({
        "index": 0, "img_paths": [input_image], "context": [prompt], "aspect_ratio": aspect_ratio
    })
    
    return _execute_vincie_logic(prompt_config, steps, cfg_scale, seed, pad_img_placeholder=True, resolution=resolution)

def run_multi_turn_inference(
    input_image: str, prompts_text: str, steps: int, cfg_scale: float, seed: int, progress=gr.Progress()
) -> List[Image.Image]:
    """Handler para a aba 'Edição em Múltiplos Turnos'."""
    if not all([input_image, prompts_text]):
        raise gr.Error("É necessário fornecer uma imagem de entrada e pelo menos um prompt.")
    
    prompts = [p.strip() for p in prompts_text.splitlines() if p.strip()]
    if not prompts:
        raise gr.Error("Nenhum prompt válido fornecido.")
        
    output_images_with_paths = []
    
    for i, prompt in enumerate(progress.tqdm(prompts, desc="Processando turnos")):
        image_paths = [input_image] + [path for _, path in output_images_with_paths]
        context_prompts = prompts[:i+1]
        
        prompt_config = OmegaConf.create({
            "index": i, "img_paths": image_paths, "context": context_prompts, "aspect_ratio": "keep_ratio"
        })
        
        turn_seed = seed if seed == -1 else seed + i
        result_image = _execute_vincie_logic(prompt_config, steps, cfg_scale, turn_seed, pad_img_placeholder=True, resolution=512)
        
        temp_path = os.path.join("/tmp", f"{uuid.uuid4()}.png")
        result_image.save(temp_path)
        output_images_with_paths.append((result_image, temp_path))

    return [img for img, _ in output_images_with_paths]

def run_multi_concept_inference(prompt: str, *images: str) -> Image.Image:
    """Handler para a aba 'Composição de Conceitos'."""
    image_paths = [img for img in images if img is not None]
    if not image_paths or not prompt.strip():
        raise gr.Error("É necessário um prompt e pelo menos uma imagem de entrada.")

    # Constrói a lista de prompts: N-1 placeholders + 1 prompt de composição.
    prompts_list = [f"<IMG{i}>: " for i in range(1, len(image_paths))]
    prompts_list.append(prompt)
        
    prompt_config = OmegaConf.create({
        "index": 0, "img_paths": image_paths, "context": prompts_list, "aspect_ratio": "1:1"
    })
    
    # Usa parâmetros fixos para esta funcionalidade, conforme documentação do VINCIE.
    return _execute_vincie_logic(prompt_config, steps=50, cfg_scale=7.5, seed=1, pad_img_placeholder=False, resolution=512)

def _log_pipeline_args(prompt_config, steps, cfg_scale, final_seed, resolution, pad_placeholder):
    """Função auxiliar para imprimir os argumentos exatos enviados à pipeline do VINCIE."""
    log_data = {
        "--- INÍCIO DOS ARGUMENTOS DA PIPELINE VINCIE ---": "",
        "1. Configuração do Prompt": OmegaConf.to_container(prompt_config, resolve=True),
        "2. Parâmetros de Difusão": {
            "steps": int(steps),
            "cfg_scale": float(cfg_scale),
        },
        "3. Parâmetros de Geração": {
            "seed_final": int(final_seed),
            "resolution": int(resolution),
            "pad_img_placeholder": bool(pad_placeholder),
        },
        "--- FIM DOS ARGUMENTOS ---": ""
    }
    print(json.dumps(log_data, indent=2, ensure_ascii=False))

# --- 4. CONSTRUÇÃO DA INTERFACE GRADIO ---

def create_ui():
    """Cria e retorna a interface Gradio completa com todas as abas e controles."""
    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="VINCIE Playground") as demo:
        gr.Markdown("# 🖼️ **VINCIE Playground**\nExplore as diferentes capacidades do modelo VINCIE.")

        # Controles avançados compartilhados pelas abas 1 e 2
        with gr.Accordion("Opções Avançadas (para Abas 1 e 2)", open=False):
            steps_input = gr.Slider(label="Passos de Inferência", minimum=10, maximum=100, step=1, value=50)
            cfg_scale_input = gr.Slider(label="Escala de Orientação (CFG)", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
            seed_input = gr.Number(label="Semente (Seed)", value=-1, precision=0, info="Use -1 para aleatório.")

        with gr.Tabs():
            # Aba 1: Edição Simples
            with gr.TabItem("Edição Simples"):
                with gr.Row(equal_height=False):
                    with gr.Column(scale=1):
                        single_turn_img_in = gr.Image(type="filepath", label="Imagem de Entrada")
                        single_turn_prompt = gr.Textbox(lines=2, label="Prompt de Edição")
                        with gr.Accordion("Opções de Imagem", open=True):
                             aspect_ratio_input = gr.Dropdown(label="Aspect Ratio", choices=["keep_ratio", "1:1", "16:9", "9:16", "4:3", "3:4"], value="keep_ratio")
                             resolution_input = gr.Slider(label="Resolução (lado menor)", minimum=256, maximum=1024, step=64, value=512)
                        single_turn_button = gr.Button("Gerar", variant="primary")
                    with gr.Column(scale=1):
                        single_turn_img_out = gr.Image(label="Resultado", interactive=False)
                gr.Examples([["/app/VINCIE/assets/woman_pineapple.png", "Adicione uma coroa na cabeça da mulher."]], [single_turn_img_in, single_turn_prompt])

            # Aba 2: Edição em Múltiplos Turnos
            with gr.TabItem("Edição em Múltiplos Turnos"):
                with gr.Row():
                    with gr.Column(scale=1):
                        multi_turn_img_in = gr.Image(type="filepath", label="Imagem de Entrada Inicial")
                        multi_turn_prompts = gr.Textbox(lines=5, label="Prompts (um por linha)", placeholder="Turno 1: faça isso\nTurno 2: agora mude aquilo...")
                        multi_turn_button = gr.Button("Gerar Sequência", variant="primary")
                    with gr.Column(scale=2):
                        multi_turn_gallery_out = gr.Gallery(label="Resultados dos Turnos", columns=3, height="auto")

            # Aba 3: Composição de Conceitos
            with gr.TabItem("Composição de Conceitos"):
                gr.Markdown("Faça o upload de até 6 imagens (`<IMG0>` a `<IMG5>`) e escreva um prompt que as combine para gerar uma nova imagem (`<IMG6>`).")
                with gr.Row():
                    concept_inputs = [gr.Image(type="filepath", label=f"Imagem {i} (<IMG{i}>)") for i in range(6)]
                concept_prompt = gr.Textbox(lines=4, label="Prompt de Composição Final", value="Baseado em <IMG0> e <IMG1>, um retrato do homem de <IMG0> usando o chapéu de <IMG1>. Saída <IMG6>:")
                concept_button = gr.Button("Compor Imagem", variant="primary")
                concept_img_out = gr.Image(label="Resultado da Composição", interactive=False)

        # Conecta os botões às suas respectivas funções de backend
        single_turn_button.click(fn=run_single_turn_inference, inputs=[single_turn_img_in, single_turn_prompt, aspect_ratio_input, resolution_input, steps_input, cfg_scale_input, seed_input], outputs=[single_turn_img_out])
        multi_turn_button.click(fn=run_multi_turn_inference, inputs=[multi_turn_img_in, multi_turn_prompts, steps_input, cfg_scale_input, seed_input], outputs=[multi_turn_gallery_out])
        concept_button.click(fn=run_multi_concept_inference, inputs=[concept_prompt] + concept_inputs, outputs=[concept_img_out])
        
    return demo

# --- 5. PONTO DE ENTRADA DA APLICAÇÃO ---

if __name__ == "__main__":
    setup_model()
    ui = create_ui()
    
    server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0")
    server_port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
    enable_queue = os.environ.get("GRADIO_ENABLE_QUEUE", "True").lower() == "true"

    print(f"INFO: Lançando a interface Gradio em http://{server_name}:{server_port}")
    if enable_queue:
        ui.queue()

    ui.launch(server_name=server_name, server_port=server_port)