# -*- coding: utf-8 -*-
"""
MÚSICA — ACE-Step 1.5 XL Turbo, no formato DIFFUSERS.

É a primeira vez que este app tem MÚSICA de verdade sem provedor pago: a
trilha do `midia_livre` é síntese de acordes (e a tela sempre disse isso);
o ACE-Step compõe — com letra, inclusive. Quando o worker está de pé, é ele
quem responde "music"; quando dorme, a trilha sintetizada volta a valer, e a
diferença de qualidade é o preço anunciado da GPU gratuita.

POR QUE DIFFUSERS E NÃO O PACOTE `acestep`, decidido na primeira validação
real: o pacote do projeto crava `transformers==4.50.0`, e o Z-Image/Qwen3-TTS
exigem >=4.51 — no mesmo processo os dois não se resolvem (pip cai em
ResolutionImpossible). O projeto publica o checkpoint OFICIAL em formato
diffusers (`acestep-v15-xl-turbo-diffusers`), carregável pelo
`AceStepPipeline` que entrou no diffusers 0.38 — e o text_encoder dele é
justamente um Qwen3Model, o mesmo stack dos outros modelos. Um requirements
só, os quatro tipos juntos. (O id antigo `ACE-Step/ACE-Step-v1.5` nem existe
no Hub — 404 medido.)

`lyrics` aceita o formato de seções do projeto ([Verse]/[Chorus]); sem letra
vai "[instrumental]" — mandar string vazia faz alguns checkpoints alucinarem
voz. O turbo é destilado: 8 passos, sem CFG para ajustar.
"""
from __future__ import annotations

import io
import logging
import os
from typing import Dict, Optional

from . import base

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

VRAM_MINIMA_GB = 12

_pipe = None
_modelo_carregado = ""


def _modelo() -> str:
    return os.getenv("MUSIC_MODEL",
                     "ACE-Step/acestep-v15-xl-turbo-diffusers").strip()


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 música…")
    import torch
    from diffusers import AceStepPipeline
    _pipe = AceStepPipeline.from_pretrained(
        _modelo(), torch_dtype=torch.bfloat16,
        token=os.getenv("HF_TOKEN") or None)
    # Tiling no VAE: decodificar minutos de áudio 48 kHz de uma vez é o pico
    # de memória do pipeline — o card do modelo recomenda para áudio longo.
    _pipe.vae.enable_tiling()
    livre = base.vram_livre_gb()
    if livre is not None and livre < 16:
        _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 base.wav_silencio(0.5), "audio/wav"

    prompt = (params.get("prompt") or "").strip()
    estilo = (params.get("style") or "").strip()
    if estilo:
        prompt = f"{estilo}, {prompt}"
    letra = (params.get("lyrics") or "").strip() or "[instrumental]"
    # Duração com teto: cada segundo é tempo de GPU, e um pedido de "3600"
    # vindo de um cliente forjado não pode virar uma hora de placa ocupada.
    duracao = max(10, min(int(params.get("duration") or 60),
                          int(os.getenv("MUSIC_MAX_SEGUNDOS", "240"))))

    progresso(5, "Compondo…")
    saida = _pipe(prompt=prompt, lyrics=letra,
                  audio_duration=float(duracao),
                  generator=base.semente(params))
    progresso(95, "Finalizando…")
    # audios[0] é (canais, amostras) em 48 kHz; o WAV sai estéreo.
    return _wav_estereo(saida.audios[0],
                        getattr(_pipe, "sample_rate", 48000)), "audio/wav"


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


def config_menor(params: Dict) -> Optional[Dict]:
    p = dict(params)
    dur = int(p.get("duration") or 60)
    if dur > 30:
        p["duration"] = 30
        return p
    return None


def _wav_estereo(onda, taxa: int) -> bytes:
    """Tensor (canais, amostras) → WAV 16-bit. struct, sem soundfile —
    uma dependência a menos, como no TTS."""
    import struct
    try:
        dados = onda.detach().float().cpu().numpy()
    except AttributeError:
        dados = onda
    if dados.ndim == 1:
        dados = dados.reshape(1, -1)
    canais = dados.shape[0]
    # Intercala os canais (LRLR…), que é como o WAV guarda estéreo.
    inteiros = (dados.T.clip(-1, 1) * 32767).astype("<i2").tobytes()
    bloco = canais * 2
    return (b"RIFF" + struct.pack("<I", 36 + len(inteiros)) + b"WAVEfmt "
            + struct.pack("<IHHIIHH", 16, 1, canais, int(taxa),
                          int(taxa) * bloco, bloco, 16)
            + b"data" + struct.pack("<I", len(inteiros)) + inteiros)