#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
from pathlib import Path
from typing import Optional, Tuple
import gradio as gr

from services.seedvr import SeedVRRefineService

APP_HOME = Path(os.environ.get("APP_HOME", "/app"))
OUT_DIR = APP_HOME / "outputs" / "seedvr_refine"

svc = SeedVRRefineService()

def setup() -> str:
    try:
        token = os.environ.get("HF_TOKEN") or os.environ.get("HFTOKEN")
        msg = svc.ensure_model(max_workers=int(os.environ.get("MAX_WORKERS", "48")), token=token)
        svc.ensure_apex(enable_shim=True)
        return f"Setup concluído: {msg} (apex shim ativo)."
    except Exception as e:
        return f"Setup falhou: {e}"

def refine_ui(
    input_file: gr.File,
    upscale: float,
    strength: float,
    denoise: float,
    t_consistency: float,
    fps_out: int,
    tile: int,
    dtype: str
) -> Tuple[str, str]:
    try:
        inp = Path(input_file.name) if hasattr(input_file, "name") else Path(input_file)
        out = svc.refine(
            input_path=inp,
            output_dir=OUT_DIR,
            upscale=upscale,
            strength=strength,
            denoise=denoise,
            t_consistency=t_consistency,
            fps_out=(None if fps_out <= 0 else fps_out),
            tile=(None if tile <= 0 else tile),
            dtype=(None if not dtype else dtype)
        )
        return (f"OK: saída em {out}", out)
    except Exception as e:
        return (f"Erro: {e}", "")

with gr.Blocks(title="SeedVR — Refinamento (sem Apex)", theme="Default") as demo:
    gr.Markdown("### SeedVR — Refinamento de vídeo/imagem (qualidade, sem Apex)")
    with gr.Row():
        with gr.Column(scale=2):
            input_file = gr.File(label="Vídeo/Imagem de entrada", type="filepath")
            upscale = gr.Slider(1.0, 2.0, value=1.0, step=0.1, label="Upscale factor")
            strength = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Força do refine")
            denoise = gr.Slider(0.0, 1.0, value=0.1, step=0.05, label="Denoise")
            tcons = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Consistência temporal")
            fps_out = gr.Number(value=0, precision=0, label="FPS saída (0 = manter)")
            tile = gr.Number(value=0, precision=0, label="Tile (0 = off)")
            dtype = gr.Dropdown(choices=["", "bfloat16", "float16"], value="", label="Dtype override")
            go = gr.Button("Refinar", variant="primary")
        with gr.Column(scale=1):
            status = gr.Textbox(label="Status")
            outdir = gr.Textbox(label="Saída", value=str(OUT_DIR), interactive=False)

    go.click(
        fn=refine_ui,
        inputs=[input_file, upscale, strength, denoise, tcons, fps_out, tile, dtype],
        outputs=[status, outdir]
    )

    init = gr.Button("Preparar modelos (setup)")
    init.click(fn=setup, outputs=status)

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    demo.launch(server_name="0.0.0.0", server_port=port, allowed_paths=[str(OUT_DIR), str(APP_HOME / "ckpt" / "SeedVR2-3B")])