# -*- coding: utf-8 -*-
"""
A CASCA GRADIO DO ADAPTER ZEROGPU — o que o Hugging Face executa.

Este arquivo é fino de propósito: toda a lógica mora em `nucleo.py` (testado
no repositório, sem Gradio). Aqui só se faz o que EXIGE o ambiente do Space:

  - decorar a geração com `@spaces.GPU(duration=...)` — é o decorator que
    faz o ZeroGPU alocar a placa durante a chamada, e a duração é POR TIPO
    (IMAGE_GPU_SECONDS etc., lidas no nucleo);
  - expor as funções como API do Gradio (`api_name=...`), que é o único
    transporte que um Space ZeroGPU oferece: POST → event_id → SSE. Os
    yields do gerador viram eventos `generating` (o progresso do app);
  - manter a PÁGINA do Space inofensiva: quem abrir a URL num navegador vê
    um aviso, não um formulário que gasta quota — gerar exige o segredo, e a
    conferência roda ANTES do decorator de GPU ser invocado.

Fora do Space (sem `spaces` instalado) o decorator vira identidade e o
WORKER_MOCK=1 permite subir a casca inteira num notebook para teste.
"""
from __future__ import annotations

import gradio as gr

import nucleo

try:
    import spaces
    _GPU = spaces.GPU
except ImportError:              # fora do Hugging Face: sem decorator
    def _GPU(*a, **k):
        def envolve(fn):
            return fn
        return envolve


def _fazer_endpoint(tipo: str):
    """Uma função por tipo, porque a DURAÇÃO da GPU é do decorator e vídeo
    não pode viver com o prazo de imagem (abortaria no meio, com a quota já
    gasta)."""

    @_GPU(duration=nucleo.duracao_gpu(tipo))
    def _na_gpu(params_json: str):
        yield from nucleo.gerar(tipo, params_json)

    def endpoint(segredo: str, params_json: str):
        # A porta fica AQUI FORA: recusar antes de `_na_gpu` significa que a
        # chamada sem segredo não aloca GPU e não consome quota de ninguém.
        recusa = nucleo.autenticar(segredo)
        if recusa:
            yield {"erro": recusa}
            return
        # O DOWNLOAD dos pesos roda AQUI FORA da GPU: é rede, não CUDA, e
        # dentro do @spaces.GPU ele queimava o prazo — a duração declarada só
        # precisa cobrir carregar-do-disco + inferência.
        yield {"etapa": "Baixando o modelo…", "progress": None}
        falha = nucleo.preparar(tipo)
        if falha:
            # erro_gpu, e não erro: download que falhou (repo errado, rede,
            # Hub fora) é ESTE runtime não conseguindo atender agora — o
            # backend degrada para o gerador local em vez de devolver a
            # falha à pessoa. Medido: o 404 de um id de modelo errado saía
            # como erro do pedido e a música morria sem o fallback.
            yield {"erro_gpu": falha}
            return
        # O try daqui não é redundância do try do nucleo: a ALOCAÇÃO do
        # @spaces.GPU acontece ANTES de o corpo rodar, e é ali que nascem os
        # erros de quota e de duração do ZeroGPU. Sem esta cerca eles
        # estouravam dentro do Gradio, que os engole (show_error desligado) —
        # o backend recebia "o Space abortou a chamada: None", sem motivo
        # nenhum para agir. Medido na primeira geração real.
        try:
            yield from _na_gpu(params_json)
        except BaseException as e:
            if isinstance(e, (KeyboardInterrupt, SystemExit)):
                raise
            if nucleo.e_erro_de_gpu(e):
                yield {"erro_gpu": f"ZeroGPU indisponível agora: {str(e)[:200]}"}
            else:
                yield {"erro": f"{type(e).__name__}: {str(e)[:300]}"}

    return endpoint


def _saude(segredo: str):
    recusa = nucleo.autenticar(segredo)
    if recusa:
        return {"erro": recusa}
    return nucleo.saude()


with gr.Blocks(title="Conecta Primo — GPU") as demo:
    # A página humana diz o mínimo. Nada aqui dispara geração: os botões nem
    # existem — só a API, e ela pede o segredo.
    gr.Markdown("## Conecta Primo — worker de GPU\n"
                "Este Space atende **somente** o backend do aplicativo, "
                "autenticado por segredo. Não há nada para usar nesta página.")

    segredo_in = gr.Textbox(visible=False)
    params_in = gr.Textbox(visible=False)
    saida = gr.JSON(visible=False)

    # gr.api registra função como endpoint puro (Gradio >= 4.44). Os botões
    # invisíveis são a forma equivalente que funciona em toda versão 4/5 —
    # ficam os dois? Não: DOIS caminhos para a mesma função é o defeito de
    # sempre. Fica o dos componentes, que é o estável.
    for _tipo, _nome in (("image", "gerar_imagem"), ("voice", "gerar_fala"),
                         ("music", "gerar_musica"), ("video", "gerar_video")):
        _btn = gr.Button(visible=False)
        _btn.click(_fazer_endpoint(_tipo), inputs=[segredo_in, params_in],
                   outputs=[saida], api_name=_nome)

    _btn_saude = gr.Button(visible=False)
    _btn_saude.click(_saude, inputs=[segredo_in], outputs=[saida],
                     api_name="saude")

if __name__ == "__main__":
    # Fila do Gradio LIGADA com concorrência 1: no ZeroGPU cada chamada é uma
    # alocação de GPU, e duas em paralelo é OOM ou quota dobrada — a mesma
    # decisão do MAX_CONCURRENT_JOBS=1 do worker CUDA, no dialeto do Gradio.
    demo.queue(default_concurrency_limit=1, max_size=16)
    demo.launch(show_api=False)