import os

# Отключаем SSR до импорта Gradio
os.environ["GRADIO_SSR_MODE"] = "false"

import gc
import tempfile

import gradio as gr
import soundfile as sf
import torch
import spaces

from diffusers import ModularPipeline


MODEL_ID = "Cartik/Sonexa-Music-v0.1-Beta"


# ============================================================
# CONFIG
# ============================================================

print("=" * 60)
print("Sonexa Music Generation")
print("=" * 60)
print(f"Model: {MODEL_ID}")
print(f"CUDA available: {torch.cuda.is_available()}")
print("=" * 60)


# ============================================================
# LOAD MODEL
# ============================================================

print("Loading music model on CPU...")

pipe = ModularPipeline.from_pretrained(
    MODEL_ID,
)

pipe.load_components(
    dtype=torch.float32,
)

# Модель постоянно находится в RAM.
# ZeroGPU выдаёт GPU только при вызове @spaces.GPU.
pipe.to("cpu")

print("Music model loaded successfully!")


# ============================================================
# MUSIC GENERATION
# ============================================================

@spaces.GPU(duration=120)
def generate_music(prompt, lyrics, duration):

    prompt = (prompt or "").strip()
    lyrics = (lyrics or "").strip()

    if not prompt:
        raise gr.Error(
            "Опиши, какую музыку нужно создать."
        )

    try:
        duration = float(duration)
        duration = max(10, min(duration, 120))

        print("=" * 60)
        print("GPU GENERATION START")
        print(f"Prompt: {prompt}")
        print(f"Lyrics: {lyrics[:200]}")
        print(f"Duration: {duration}s")
        print("=" * 60)

        # ZeroGPU выделяет CUDA перед входом сюда.
        if not torch.cuda.is_available():
            raise RuntimeError(
                "ZeroGPU не выделил CUDA для генерации."
            )

        device = "cuda"

        print("Moving music model to GPU...")

        pipe.to(device)

        print("Model moved to GPU.")

        generator = torch.Generator(
            device=device
        ).manual_seed(
            int.from_bytes(
                os.urandom(4),
                "little",
            )
        )

        print("Starting inference...")

        with torch.inference_mode():

            result = pipe(
                prompt=prompt,
                lyrics=lyrics,
                audio_duration=duration,
                generator=generator,
                output="audios",
            )

        print("Inference finished.")

        audio = result[0]

        # ====================================================
        # SAVE AUDIO
        # ====================================================

        output_file = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False,
        )

        output_path = output_file.name
        output_file.close()

        # Модель возвращает channel-first audio.
        audio_array = (
            audio
            .T
            .float()
            .cpu()
            .numpy()
        )

        sf.write(
            output_path,
            audio_array,
            pipe.sampling_rate,
        )

        print(
            f"Music saved: {output_path}"
        )

        # Возвращаем модель в RAM.
        print("Moving model back to CPU...")

        pipe.to("cpu")

        print("Model returned to CPU.")

        return output_path

    except Exception as e:

        print("=" * 60)
        print("MUSIC GENERATION ERROR")
        print("=" * 60)
        print(repr(e))
        print("=" * 60)

        # Пытаемся вернуть модель в RAM,
        # даже если генерация упала.
        try:
            pipe.to("cpu")
        except Exception as cpu_error:
            print(
                f"Could not move model to CPU: {cpu_error}"
            )

        raise gr.Error(
            f"Ошибка генерации: {e}"
        )

    finally:

        gc.collect()

        if torch.cuda.is_available():
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass

        print("Generation cleanup finished.")


# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title="Sonexa Music"
) as demo:

    gr.Markdown(
        "# Sonexa Music Generation"
    )

    gr.Markdown(
        "Создавайте музыку по описанию и тексту песни."
    )

    with gr.Row():

        # ====================================================
        # INPUT
        # ====================================================

        with gr.Column():

            prompt = gr.Textbox(
                label="Описание музыки",
                placeholder=(
                    "Energetic electronic music, "
                    "powerful drums, deep bass, "
                    "bright synths and an intense drop"
                ),
                lines=5,
            )

            lyrics = gr.Textbox(
                label="Текст песни",
                placeholder=(
                    "[Verse]\n"
                    "Morning light is falling...\n\n"
                    "[Chorus]\n"
                    "We keep moving through the night..."
                ),
                lines=8,
            )

            duration = gr.Slider(
                minimum=10,
                maximum=120,
                value=60,
                step=10,
                label="Длительность (секунды)",
            )

            button = gr.Button(
                "Создать музыку",
                variant="primary",
            )

        # ====================================================
        # OUTPUT
        # ====================================================

        with gr.Column():

            output = gr.Audio(
                label="Результат",
                type="filepath",
            )


    # ========================================================
    # API
    # ========================================================

    button.click(
        fn=generate_music,
        inputs=[
            prompt,
            lyrics,
            duration,
        ],
        outputs=output,
        api_name="predict",
    )


# ============================================================
# QUEUE
# ============================================================

demo.queue(
    max_size=8,
    default_concurrency_limit=1,
)


# ============================================================
# START
# ============================================================

print("=" * 60)
print("Starting Gradio...")
print("SSR disabled through environment.")
print("=" * 60)

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_api=True,
)