# -*- coding: utf-8 -*-
"""
VÍDEO — Wan 2.2 (TI2V 5B), texto-para-vídeo e imagem-para-vídeo.

A variante 5B é a escolhida de propósito: é a que o próprio projeto publica
como "cabe em uma GPU de consumidor" (24 GB com offload). As variantes A14B
são melhores e NÃO cabem no worker que este projeto consegue pagar — quem
tiver um A100 sobrando troca o VIDEO_MODEL e nada mais.

As duas rotas do backend caem AQUI: `image_url` presente = image-to-video
(o primeiro quadro é a imagem da pessoa); ausente = text-to-video. Um módulo
só, porque é o mesmo pipeline com um argumento a mais — dois módulos seriam
duas cópias do mesmo offload para divergir.

── POR QUE O PADRÃO NÃO É 720P, E ISSO É CONTA E NÃO GOSTO ─────────────────
O prazo de GPU desta casa é de 180 SEGUNDOS, e ele é o maior valor que o
ZeroGPU gratuito aceita (ver CLAUDE.md: 240 declarados viraram "larger than
the maximum allowed"). O cartão oficial deste modelo diz, com estas palavras,
que ele gera "a 5-second 720P video in under 9 minutes on a single
consumer-grade GPU" — 540 segundos, em 50 passos, a 1280x704.

O prazo é DURO e não tem rede de segurança: `config_menor` só é chamado
depois de um OOM, e estourar o TEMPO não é OOM — é a chamada abortada no
meio, com a quota do dia gasta e nada saindo. Ou seja, o padrão de 720p era
um vídeo que praticamente nunca terminava, e o sintoma disso não é um erro
que aponte para cá: é a geração degradando para o gerador de reserva, como se
a GPU estivesse dormindo.

Então o padrão é 832x480 — a MESMA medida que o `config_menor` abaixo já
escolhia como primeiro degrau, e portanto nenhuma resolução nova entrando no
modelo por esta porta. São 2,3x menos pixels. Com `GPU_DEDICADA=1` volta o
1280x704, para o dia em que houver prazo de verdade.

Vídeo 480p que sai é infinitamente melhor que 720p que aborta.

── E A PROPORÇÃO É DE QUEM PEDIU, não de quem escreveu este arquivo ──────────
Relato do dono: "o vídeo sai sempre horizontal". A causa não era o modelo — o
Wan gera vertical tão bem quanto horizontal —, era o `_payload_para` do
backend escrevendo 1280x720 À MÃO em todo pedido e jogando fora a proporção
que a pessoa tinha escolhido. Agora o payload leva `aspect_ratio`, e a tabela
que a traduz em pixels mora AQUI, em `_medida`, num lugar só: os dois lados
da tabela (o que cabe no prazo e o que a GPU dedicada aguenta) têm a MESMA
contagem de pixels por linha, então trocar de proporção nunca troca de risco
de estourar o relógio. Toda medida é múltiplo de 16 — o VAE do TI2V-5B
comprime 16x16, e uma medida quebrada é erro de shape no meio da geração.
"""
from __future__ import annotations

import base64 as b64
import io
import logging
import os
import tempfile
from typing import Dict, Optional

from . import base

logger = logging.getLogger("worker.wan22")

VRAM_MINIMA_GB = 16

_pipe = None
_modelo_carregado = ""


def _modelo() -> str:
    return os.getenv("VIDEO_MODEL", "Wan-AI/Wan2.2-TI2V-5B-Diffusers").strip()


def _quer_720p() -> bool:
    """O 720p nativo do cartão só com prazo de GPU de verdade.

    Quem responde é `base.gpu_dedicada()`, e não uma variável própria daqui:
    a resolução do vídeo e o número de passos da imagem nascem do MESMO fato
    (o teto de 180 s do ZeroGPU gratuito), e duas variáveis para um fato só
    é uma delas ficando para trás no dia da GPU dedicada.
    """
    return base.gpu_dedicada()


# Proporção -> (largura, altura). Duas tabelas com a MESMA área por linha (ver
# o cabeçalho): ~400 mil pixels na que cabe no prazo, ~900 mil na dedicada.
# Tudo múltiplo de 16.
_MEDIDA_NO_PRAZO = {"16:9": (832, 480), "9:16": (480, 832), "1:1": (624, 624)}
_MEDIDA_DEDICADA = {"16:9": (1280, 704), "9:16": (704, 1280), "1:1": (960, 960)}
ASPECTO_PADRAO = "16:9"


def _aspecto(params: Dict) -> str:
    a = str(params.get("aspect_ratio") or "").strip()
    return a if a in _MEDIDA_NO_PRAZO else ASPECTO_PADRAO


def _medida(params: Dict) -> tuple:
    """A medida em vigor, num lugar SÓ.

    Ela é lida em dois lugares (o `gerar` e o `config_menor`), e o padrão
    escrito à mão nos dois é como um deles fica para trás: com o padrão antigo
    de 1280 no `config_menor`, uma geração já em 832x480 seria "rebaixada"
    para 832x480 de novo e a escada do OOM perderia um degrau sem dizer nada.

    Quem manda `width`/`height` explícitos ganha da proporção — é o caminho
    do `config_menor`, que já escolheu a medida do degrau.
    """
    tabela = _MEDIDA_DEDICADA if _quer_720p() else _MEDIDA_NO_PRAZO
    padrao = tabela[_aspecto(params)]
    return (int(params.get("width") or padrao[0]),
            int(params.get("height") or padrao[1]))


def baixar(progresso: base.Progresso) -> None:
    base.baixar_repo(_modelo(), progresso)


def carregar(progresso: base.Progresso) -> None:
    global _pipe, _modelo_carregado
    if base.modo_mock():
        _pipe = "mock"
        return
    if _pipe is not None and _modelo_carregado == _modelo():
        return
    progresso(None, "Carregando o modelo de vídeo…")
    import torch
    from diffusers import WanPipeline
    _pipe = WanPipeline.from_pretrained(_modelo(), torch_dtype=torch.bfloat16,
                                        token=os.getenv("HF_TOKEN") or None)
    # O OFFLOAD DEIXOU DE SER INCONDICIONAL, e isto é a metade do parágrafo
    # do prazo lá em cima. Ele estava ligado SEMPRE, com o argumento de que
    # "o pico do VAE de vídeo no decode é traiçoeiro e é onde o OOM mora" —
    # e o argumento continua verdadeiro numa placa apertada. O que ele não
    # via é o preço: o offload sobe os blocos da RAM a CADA passo, e num
    # prazo de 180 s isso é a diferença entre terminar e ser abortado.
    #
    # A regra é a mesma que o `qwen_image.py` já usa, e por isso mora aqui do
    # mesmo jeito: quem decide é a VRAM que existe, não um palpite escrito à
    # mão. Com placa folgada (o H200 do ZeroGPU tem de sobra para um 5B em
    # bf16) o modelo vai inteiro para a GPU; com placa apertada o offload
    # entra como sempre entrou, e o `config_menor` continua sendo a rede de
    # segurança do OOM.
    livre = base.vram_livre_gb()
    if livre is not None and livre < 24:
        _pipe.enable_model_cpu_offload()
    else:
        _pipe.to(base.dispositivo())
    _modelo_carregado = _modelo()


def gerar(params: Dict, progresso: base.Progresso) -> base.Saida:
    if base.modo_mock():
        return _mp4_mock(), "video/mp4"

    prompt = (params.get("prompt") or "").strip()
    estilo = (params.get("style") or "").strip()
    if estilo:
        prompt = f"{prompt}, {estilo} style"
    # Quem manda width/height explícito ganha; senão a PROPORÇÃO do pedido
    # escolhe na tabela, e a tabela é a que cabe no prazo (ver o cabeçalho).
    # `GPU_DEDICADA=1` troca para a tabela nativa do cartão.
    largura, altura = _medida(params)
    # FORA DA GPU DEDICADA, O CLIPE É CURTO E OS PASSOS SÃO MENOS. O prazo do
    # vídeo no ZeroGPU caiu para 80 s (ver `nucleo._DURACAO_PADRAO`: 180
    # declarados eram recusados pela quota por chamador), e 5 s a 30 passos
    # não terminam nisso. 3 s a 20 passos é um vídeo DE VERDADE que termina —
    # um de 5 s que aborta não entrega nada.
    curto = not base.gpu_dedicada()
    segundos = max(2, min(int(params.get("duration") or (3 if curto else 5)),
                          int(os.getenv("VIDEO_MAX_SEGUNDOS", "3" if curto else "8"))))
    fps = 24
    quadros = segundos * fps + 1
    passos = int(params.get("steps") or (20 if curto else 30))

    imagem = None
    if params.get("image_url"):
        imagem = _abrir_imagem(params["image_url"], largura, altura)

    def a_cada_passo(pipe, passo, t, kw):
        progresso(int(95 * (passo + 1) / max(1, passos)), "Gerando o vídeo…")
        return kw

    extras = {"image": imagem} if imagem is not None else {}
    saida = _pipe(prompt=prompt,
                  negative_prompt=params.get("negative_prompt") or None,
                  width=largura, height=altura, num_frames=quadros,
                  num_inference_steps=passos,
                  generator=base.semente(params),
                  callback_on_step_end=a_cada_passo, **extras)

    progresso(97, "Codificando o MP4…")
    from diffusers.utils import export_to_video
    with tempfile.TemporaryDirectory() as pasta:
        caminho = os.path.join(pasta, "video.mp4")
        export_to_video(saida.frames[0], caminho, fps=fps)
        with open(caminho, "rb") as f:
            return f.read(), "video/mp4"


def descarregar() -> None:
    global _pipe
    _pipe = None
    base.limpar_vram()


def config_menor(params: Dict) -> Optional[Dict]:
    """OOM no vídeo: primeiro menos pixels, depois menos segundos. É a ordem
    que preserva o que a pessoa pediu — um vídeo menor da cena inteira vale
    mais que meio vídeo em alta."""
    p = dict(params)
    w, h = _medida(p)
    menor = _MEDIDA_NO_PRAZO[_aspecto(p)]
    # O degrau é a medida que cabe no prazo NA MESMA PROPORÇÃO: rebaixar um
    # vídeo vertical para 832x480 entregaria um horizontal a quem pediu
    # vertical — a proporção trocada por um OOM, sem erro em lugar nenhum.
    if w * h > menor[0] * menor[1]:
        p["width"], p["height"] = menor
        return p
    dur = int(p.get("duration") or 5)
    if dur > 3:
        p["duration"] = 3
        return p
    return None


def _abrir_imagem(ref: str, largura: int, altura: int):
    """A referência chega como data URL (do app) ou como endereço http. O
    endereço http só é aceito vindo do BACKEND — e o backend já o passou pelo
    `endereco_seguro` dele; ainda assim o download daqui tem teto de tamanho,
    porque worker também não baixa gigabyte de ninguém."""
    from PIL import Image
    if ref.startswith("data:"):
        dados = b64.b64decode(ref.split(",", 1)[1])
    else:
        import requests
        r = requests.get(ref, timeout=30, stream=True)
        r.raise_for_status()
        dados = r.raw.read(20 * 1024 * 1024 + 1)
        if len(dados) > 20 * 1024 * 1024:
            raise ValueError("imagem de referência grande demais")
    img = Image.open(io.BytesIO(dados)).convert("RGB")
    return img.resize((largura, altura))


def _mp4_mock() -> bytes:
    # Um MP4 mínimo (só o ftyp) — suficiente para os testes conferirem mime e
    # transporte. O modo mock nunca finge ser vídeo de verdade: o metadado do
    # job diz "mock".
    return bytes.fromhex("000000186674797069736f6d0000020069736f6d69736f32")