#!/usr/bin/env python3
import os
import glob
import gradio as gr
from pathlib import Path
from typing import List, Tuple, Optional

from vincie_service import VincieService

svc = VincieService()  # usa /app/VINCIE e /app/ckpt/VINCIE-3B por padrão

def setup() -> str:
    """
    Prepara repositório e modelos (idempotente).
    """
    try:
        svc.ensure_repo()
        svc.ensure_model()
    except Exception as e:
        return f"Setup falhou: {e}"
    return "Setup OK"

def _list_media(out_dir: Path, max_images: int = 24) -> Tuple[List[str], Optional[str]]:
    """
    Varre o diretório de saída por imagens e vídeo, retornando caminhos ordenados.
    """
    exts_img = ("*.png", "*.jpg", "*.jpeg", "*.webp")
    imgs: List[Path] = []
    for pat in exts_img:
        imgs += list(out_dir.rglob(pat))
    imgs = sorted(imgs, key=lambda p: p.stat().st_mtime)
    imgs_paths = [str(p) for p in imgs[-max_images:]] if imgs else []

    vids = sorted(out_dir.rglob("*.mp4"), key=lambda p: p.stat().st_mtime)
    vid_path = str(vids[-1]) if vids else None
    return imgs_paths, vid_path

def ui_multi_turn(input_image, turns_text):
    """
    Executa multi-turn editing e retorna mídia do diretório de saída.
    """
    if not input_image or not str(input_image).strip():
        return [], None, "Forneça uma imagem."
    if not turns_text or not turns_text.strip():
        return [], None, "Forneça turns (um por linha)."

    turns = [ln.strip() for ln in turns_text.splitlines() if ln.strip()]
    try:
        out_dir = svc.multi_turn_edit(input_image, turns)
    except Exception as e:
        return [], None, f"Erro na geração: {e}"

    imgs, vid = _list_media(Path(out_dir))
    status = f"Gerado em: {str(out_dir)}"
    return imgs, vid, status

def ui_multi_concept(files, descs_text, final_prompt):
    """
    Executa composição multi-conceito e retorna mídia do diretório de saída.
    """
    if not files:
        return [], None, "Envie imagens de conceito."
    if not descs_text or not descs_text.strip():
        return [], None, "Forneça descrições (uma por linha)."
    if not final_prompt or not final_prompt.strip():
        return [], None, "Forneça o prompt final."

    descs = [ln.strip() for ln in descs_text.splitlines() if ln.strip()]
    if len(descs) != len(files):
        return [], None, f"Número de descrições ({len(descs)}) difere do número de imagens ({len(files)})."

    try:
        out_dir = svc.multi_concept_compose(files, descs, final_prompt)
    except Exception as e:
        return [], None, f"Erro na geração: {e}"

    imgs, vid = _list_media(Path(out_dir))
    status = f"Gerado em: {str(out_dir)}"
    return imgs, vid, status

with gr.Blocks(title="VINCIE Service") as demo:
    gr.Markdown("# 🎨 VINCIE Service — Multi-turn e Multi-concept")

    with gr.Row():
        setup_btn = gr.Button("Preparar (repo + modelos)")
        setup_out = gr.Textbox(label="Status", interactive=False)
    setup_btn.click(fn=setup, outputs=setup_out)

    with gr.Tab("🔄 Multi-turn Editing"):
        with gr.Row():
            img = gr.Image(type="filepath", label="Imagem inicial")
            turns = gr.Textbox(lines=8, label="Turns (um por linha)")
        run1 = gr.Button("Executar")
        out_gallery = gr.Gallery(label="Imagens", columns=4, height="auto")
        out_video = gr.Video(label="Vídeo (se houver)")
        out_status = gr.Textbox(label="Saída", interactive=False)
        run1.click(ui_multi_turn, inputs=[img, turns], outputs=[out_gallery, out_video, out_status])

    with gr.Tab("🧩 Multi-concept Compose"):
        files = gr.File(file_count="multiple", type="filepath", label="Imagens de conceito")
        descs = gr.Textbox(lines=8, label="Descrições (uma por linha, na mesma ordem das imagens)")
        final_prompt = gr.Textbox(lines=2, label="Prompt final")
        run2 = gr.Button("Executar")
        out_gallery2 = gr.Gallery(label="Imagens", columns=4, height="auto")
        out_video2 = gr.Video(label="Vídeo (se houver)")
        out_status2 = gr.Textbox(label="Saída", interactive=False)
        run2.click(ui_multi_concept, inputs=[files, descs, final_prompt], outputs=[out_gallery2, out_video2, out_status2])

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        allowed_paths=["/app/outputs", "/app/ckpt"]
    )