"""ACE-Step 1.5 engine for Cowboy Beatz Song Lab: one clean API endpoint (/generate) on ZeroGPU.
Model init mirrors the official ACE-Step/Ace-Step-v1.5 Space (MIT)."""
import os, sys, json, time, tempfile

here = os.path.dirname(os.path.abspath(__file__))
nano = os.path.join(here, "acestep", "third_parts", "nano-vllm")
if os.path.exists(nano):
    sys.path.insert(0, nano)
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

import spaces  # before torch
import torch
import gradio as gr
from acestep.handler import AceStepHandler
from acestep.llm_inference import LLMHandler
from acestep.inference import GenerationParams, GenerationConfig, generate_music

DIT_MODEL = os.environ.get("ENGINE_DIT_MODEL", "acestep-v15-xl-turbo")
LM_MODEL = os.environ.get("ENGINE_LM_MODEL", "acestep-5Hz-lm-1.7B")
store = os.path.join(here, "data"); os.makedirs(store, exist_ok=True)

# ZeroGPU now runs on Blackwell (sm_120); the prebuilt flash-attn / flash-attn3 kernels have no sm_120
# image ("no kernel image is available"), so force PyTorch SDPA everywhere.
LLMHandler.get_best_attn_implementation = lambda self: "sdpa"
AceStepHandler.is_flash_attention_available = lambda self: False

DIT = AceStepHandler(persistent_storage_path=store)
LLM = LLMHandler(persistent_storage_path=store)
st, ok = DIT.initialize_service(project_root=here, config_path=DIT_MODEL, device="auto",
                                use_flash_attention=DIT.is_flash_attention_available(), compile_model=False,
                                offload_to_cpu=False, offload_dit_to_cpu=False)
print("DiT:", ok, st, flush=True)
lst, lok = LLM.initialize(checkpoint_dir=DIT._get_checkpoint_dir(), lm_model_path=LM_MODEL, backend="pt",
                          device="auto", offload_to_cpu=False, dtype=DIT.dtype)
print("LM:", lok, lst, flush=True)


def _dur(caption, lyrics, instrumental, bpm, keyscale, duration, seed, negative, thinking, steps):
    return int(min(170, 40 + float(duration) * 0.5))


@spaces.GPU(duration=_dur)
def generate(caption, lyrics, instrumental, bpm, keyscale, duration, seed, negative, thinking, steps):
    t0 = time.time()
    seed = int(seed) if seed is not None else -1
    instrumental = bool(instrumental)
    params = GenerationParams(
        caption=caption or "", lyrics="[Instrumental]" if instrumental else (lyrics or ""),
        instrumental=instrumental, vocal_language="unknown" if instrumental else "en",
        bpm=int(bpm) if bpm else None, keyscale=keyscale or "", timesignature="4",
        duration=float(duration), inference_steps=int(steps or 8), seed=seed,
        thinking=bool(thinking), lm_negative_prompt=negative or "NO USER INPUT",
        use_cot_metas=False, use_cot_caption=False, use_cot_language=False,
    )
    cfg = GenerationConfig(batch_size=1, use_random_seed=seed < 0, seeds=[seed] if seed >= 0 else None,
                           audio_format="wav")
    res = generate_music(DIT, LLM, params, cfg, save_dir=tempfile.mkdtemp())
    if not res.success or not res.audios or not res.audios[0].get("path"):
        raise gr.Error("Generation failed: " + str(res.error or res.status_message)[:300])
    a = res.audios[0]
    info = {"seed": a["params"].get("seed"), "gpu_seconds": round(time.time() - t0, 1),
            "time_costs": res.extra_outputs.get("time_costs", {})}
    return a["path"], json.dumps(info, default=str)


with gr.Blocks(title="ACE-Step engine (Cowboy Beatz)") as demo:
    gr.Markdown("ACE-Step 1.5 engine for the Cowboy Beatz Song Lab. API only.")
    cap = gr.Textbox(label="caption"); lyr = gr.Textbox(label="lyrics"); ins = gr.Checkbox(True, label="instrumental")
    bpm = gr.Number(150, label="bpm"); key = gr.Textbox("F minor", label="key"); dur = gr.Number(60, label="duration")
    seed = gr.Number(-1, label="seed"); neg = gr.Textbox(label="negative"); think = gr.Checkbox(True, label="thinking")
    steps = gr.Number(8, label="steps")
    out = gr.Audio(type="filepath"); meta = gr.Textbox(label="info")
    gr.Button("Generate").click(generate, [cap, lyr, ins, bpm, key, dur, seed, neg, think, steps], [out, meta], api_name="generate")

demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860, show_error=True)