import os

os.environ.setdefault("HF_HOME", "/tmp/asf-hf-cache")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

try:
    import spaces
    GPU = spaces.GPU
except Exception:
    class _GPU:
        def __init__(self, duration=60, size="large"):
            self.duration = duration
            self.size = size
        def __call__(self, fn):
            return fn
    GPU = _GPU

import gradio as gr
import torch
import soundfile as sf
import numpy as np
from diffusers import AceStepPipeline

# ASF v198.26.20 writable Hugging Face cache guard.
_ASF_HF_CACHE_ROOT = os.environ.get("ASF_HF_CACHE_DIR") or "/tmp/asf-hf-cache"
os.makedirs(_ASF_HF_CACHE_ROOT, exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "hub"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "transformers"), exist_ok=True)
os.makedirs(os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"), exist_ok=True)
os.environ.setdefault("HF_HOME", _ASF_HF_CACHE_ROOT)
os.environ.setdefault("HF_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "hub"))
os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "transformers"))
os.environ.setdefault("DIFFUSERS_CACHE", os.path.join(_ASF_HF_CACHE_ROOT, "diffusers"))

MODEL_ID = "ACE-Step/acestep-v15-xl-turbo-diffusers"
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

model_ready = False
pipeline_ready = False
last_error = ""
pipe = None

try:
    pipe = AceStepPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN,
    )
    pipe = pipe.to("cuda")
    pipe.vae.enable_tiling()
    model_ready = True
    pipeline_ready = True
except Exception as e:
    last_error = f"{type(e).__name__}: {e}"
    pipe = None


def health():
    return {
        "status": "ok" if pipeline_ready else "error",
        "model_ready": model_ready,
        "pipeline_ready": pipeline_ready,
        "model_family": "diffusers_full_pipeline",
        "loader_strategy": "DiffusionPipeline_or_family_specific_from_pretrained",
        "last_error": last_error,
        "expected_output_type": "audio",
    }


@spaces.GPU(duration=120)
@GPU(duration=120)
def generate(prompt, lyrics, audio_duration, seed):
    if not pipeline_ready or pipe is None:
        raise RuntimeError(f"Pipeline is not ready. Last error: {last_error}")
    try:
        generator = None
        if seed is not None and int(seed) >= 0:
            generator = torch.Generator(device="cuda").manual_seed(int(seed))
        output = pipe(
            prompt=prompt,
            lyrics=lyrics or "",
            audio_duration=float(audio_duration),
            generator=generator,
        )
        audio = output.audios[0]  # Tensor shape (channels, samples) or (samples,)
        sr = getattr(pipe, "sample_rate", 48000)
        audio_np = audio.cpu().float().numpy()
        # Ensure shape compatible with soundfile: (samples,) or (samples, channels)
        if audio_np.ndim == 2 and audio_np.shape[0] in (1, 2):
            audio_np = audio_np.T
        tmp_path = os.path.join("/tmp", f"acestep_{os.urandom(4).hex()}.wav")
        sf.write(tmp_path, audio_np, sr)
        return tmp_path
    except Exception as e:
        # Surface concrete error
        raise RuntimeError(f"Generation failed: {type(e).__name__}: {e}")


with gr.Blocks(title="ACE-Step v1.5 XL Turbo") as demo:
    gr.Markdown("# ACE-Step v1.5 XL Turbo — Text-to-Music")
    gr.Markdown(
        "Generate music from a text prompt. Optional lyrics can guide the generation. "
        "This model runs locally with 8-step flow-matching inference."
    )

    with gr.Tab("Generate"):
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the music you want...",
                lines=2,
            )
            lyrics = gr.Textbox(
                label="Lyrics (optional)",
                placeholder="[Verse]\nNeon lights are calling me\n[Chorus]\nRide the wave tonight",
                lines=4,
            )
        with gr.Row():
            duration = gr.Slider(
                minimum=5,
                maximum=120,
                value=30,
                step=1,
                label="Duration (seconds)",
            )
            seed = gr.Number(
                value=42,
                precision=0,
                label="Seed (-1 for random)",
            )
        gen_btn = gr.Button("Generate")
        audio_out = gr.Audio(label="Generated Audio", type="filepath", format="wav")

        gen_btn.click(
            generate,
            inputs=[prompt, lyrics, duration, seed],
            outputs=audio_out,
            api_name="/generate",
        )

        gr.Examples(
            examples=[
                ["An upbeat synthwave track with driving drums and a catchy lead", "", 30, 42],
                ["A soft piano ballad with gentle strings", "Stars above us\nQuiet nights\nForever here", 20, 7],
                ["Energetic drum and bass with deep basslines", "", 15, 99],
            ],
            inputs=[prompt, lyrics, duration, seed],
            label="Examples",
        )

    with gr.Tab("Health"):
        health_json = gr.JSON(label="Health")
        health_btn = gr.Button("Check Health", visible=False)
        health_btn.click(
            health,
            inputs=None,
            outputs=health_json,
            api_name="health",
        )

if __name__ == "__main__":
    demo.launch(show_error=True)