"""
ACE Studio — HuggingFace Space Entry Point (gr.Server mode)
============================================================
یک تجربه‌ی یک‌جعبه‌ای («توصیف کن → مدل ما آهنگ می‌سازد»)، ساخته‌شده روی همان
هسته‌ی AceStepHandler خودمان اما با معماری gr.Server (به‌جای Blocks کلاسیک):
یک صفحه‌ی HTML کاملاً اختصاصی (index.html) صفحه‌ی اصلی است و از طریق
@gradio/client مستقیماً با endpointهای API ما (که با @spaces.GPU مشخص شده‌اند)
صحبت می‌کند.

نکته‌ی حیاتی برای ZeroGPU: هنوز هم دقیقاً از demo.launch() استفاده می‌کنیم
(چون Server.launch() داخلش خودش یک Blocks می‌سازد و blocks.launch() را صدا
می‌زند — همان متدی که پکیج spaces پچ می‌کند تا توابع @spaces.GPU را در
استارتاپ به ZeroGPU گزارش بدهد).
"""
import os
import sys
import re
import json
import base64
import tempfile
import traceback

# Get current directory (app.py location)
current_dir = os.path.dirname(os.path.abspath(__file__))

# Add nano-vllm to Python path (local package)
nano_vllm_path = os.path.join(current_dir, "acestep", "third_parts", "nano-vllm")
if os.path.exists(nano_vllm_path):
    sys.path.insert(0, nano_vllm_path)

# Disable Gradio analytics
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

# Clear proxy settings that may affect Gradio
for proxy_var in ['http_proxy', 'https_proxy', 'HTTP_PROXY', 'HTTPS_PROXY', 'ALL_PROXY']:
    os.environ.pop(proxy_var, None)

# Import spaces for ZeroGPU support (must be imported before torch for proper interception)
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

import numpy as np
import soundfile as sf
import torch
from acestep.handler import AceStepHandler

# اعمال پچ اصلاحی برای غیرفعال کردن Flash Attention ناسازگار روی ZeroGPU و استفاده از موتور پایدار sdpa
AceStepHandler.is_flash_attention_available = lambda self: False
AceStepHandler.is_flash_attn3_available = lambda self: False
AceStepHandler.get_best_attn_implementation = lambda self: "sdpa"

# Detect ZeroGPU environment
IS_HUGGINGFACE_SPACE = os.environ.get("SPACE_ID") is not None
IS_ZEROGPU = IS_HUGGINGFACE_SPACE or os.environ.get("ZEROGPU") is not None


def get_persistent_storage_path():
    """Detect and return a writable persistent storage path."""
    hf_data_path = "/data"
    if os.path.exists(hf_data_path):
        try:
            test_file = os.path.join(hf_data_path, ".write_test")
            with open(test_file, 'w') as f:
                f.write("test")
            os.remove(test_file)
            print(f"Using HuggingFace persistent storage: {hf_data_path}")
            return hf_data_path
        except (PermissionError, OSError) as e:
            print(f"Warning: /data exists but is not writable: {e}")

    fallback_path = os.path.join(current_dir, "data")
    os.makedirs(fallback_path, exist_ok=True)
    print(f"Using local storage (non-persistent): {fallback_path}")
    return fallback_path


# ── Model Loading (our own high-speed / turbo checkpoint) ────────────────────
print("=" * 60)
print("ACE Studio starting up")
if IS_ZEROGPU:
    print("ZeroGPU environment detected — GPU allocated on-demand")
print("=" * 60)

_storage = get_persistent_storage_path()
handler = AceStepHandler(persistent_storage_path=_storage)

# مدل جدید و پرسرعت ما: acestep-v15-xl-turbo (۸ استپ، بدون CFG، تولید در چند ثانیه).
# قابل override با متغیر محیطی SERVICE_MODE_DIT_MODEL در صورت نیاز به مدل کیفیت بالاتر.
DIT_MODEL = os.environ.get("SERVICE_MODE_DIT_MODEL", "acestep-v15-xl-turbo")

print(f"Initializing DiT model: {DIT_MODEL}...")
_status, _ready = handler.initialize_service(
    project_root=current_dir,
    config_path=DIT_MODEL,
    device="auto",
    use_flash_attention=handler.is_flash_attention_available(),
    compile_model=False,
    offload_to_cpu=False,
    offload_dit_to_cpu=False,
)
print(f"Handler ready={_ready} — {_status}")


# ── LLM Compose (description → title / tags / lyrics) ───────────────────────
COMPOSE_SYSTEM = """You are a Grammy-winning songwriter and music producer. The user will describe a song idea in plain English. Your job is to flesh it out into a complete song specification.

Return EXACTLY this format — no extra text:

---
title: <short catchy song title>
tags: <genre and style tags, comma-separated, 3-6 tags>
bpm: <tempo as integer>
language: <vocal language: en, zh, ja, ko, or "unknown" for instrumental>
---

<song lyrics with [Verse], [Chorus], [Bridge] markers>
<use [Instrumental] alone if the song has no vocals>"""


def _compose_fallback(description: str) -> dict:
    """
    No-LLM fallback used when HF_TOKEN isn't configured (or the LLM call fails).
    We can't have an AI write lyrics without a token, but ACE-Step can still
    generate a real instrumental/style track directly from the description as
    a caption, so the app keeps working instead of hard-failing.
    """
    text = (description or "").strip()
    title = " ".join(text.split()[:6]).title() or "Untitled"
    tags = text[:200] if text else "ambient, instrumental"
    return {"title": title, "tags": tags, "lyrics": "[Instrumental]", "bpm": 120, "language": "unknown"}


def _compose(description: str) -> dict:
    """Call an LLM (via HF Inference Router) to generate tags + lyrics from a description.
    Falls back to a no-LLM heuristic if HF_TOKEN is missing or the call fails, so the
    Space still produces a track instead of erroring out."""
    key = os.environ.get("HF_TOKEN", "")
    if not key:
        print("[compose] HF_TOKEN not configured — using no-LLM fallback (instrumental)")
        return _compose_fallback(description)

    try:
        from openai import OpenAI
        client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=key)
        resp = client.chat.completions.create(
            model="openai/gpt-oss-120b:groq",
            messages=[
                {"role": "system", "content": COMPOSE_SYSTEM},
                {"role": "user", "content": description},
            ],
            max_tokens=2000,
            temperature=0.9,
        )
    except Exception as e:
        print(f"[compose] LLM call failed ({e}) — using no-LLM fallback (instrumental)")
        return _compose_fallback(description)

    raw = resp.choices[0].message.content or ""
    content = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()

    title, tags, bpm, language = "Untitled", "", 120, "en"
    lyrics = content
    m = re.search(r"---\s*\n(.*?)\n---\s*\n(.*)", content, re.DOTALL)
    if m:
        header, lyrics = m.group(1), m.group(2).strip()
        for line in header.strip().split("\n"):
            if line.startswith("title:"):
                title = line[6:].strip().strip('"\'')
            elif line.startswith("tags:"):
                tags = line[5:].strip()
            elif line.startswith("bpm:"):
                try:
                    bpm = int(line[4:].strip())
                except ValueError:
                    pass
            elif line.startswith("language:"):
                language = line[9:].strip()

    return {"title": title, "tags": tags, "lyrics": lyrics, "bpm": bpm, "language": language}


# ── GPU Inference ─────────────────────────────────────────────────────────────

def _run_inference(prompt, lyrics, audio_duration, infer_steps, seed) -> str:
    """Core inference using our AceStepHandler. Returns path to saved WAV."""
    use_random = seed < 0
    result = handler.generate_music(
        captions=prompt,
        lyrics=lyrics,
        audio_duration=audio_duration,
        inference_steps=infer_steps,
        guidance_scale=7.0,
        use_random_seed=use_random,
        seed=None if use_random else seed,
        infer_method="ode",
        shift=1.0,
        use_adg=False,
        vocal_language="en",
    )

    if not result.get("success"):
        raise RuntimeError(result.get("error", "generation failed"))

    audio_dict = result["audios"][0]
    tensor = audio_dict["tensor"]
    sr = audio_dict["sample_rate"]

    data = tensor.cpu().float().numpy()
    if data.ndim == 2:
        data = data.T
        if data.shape[1] == 1:
            data = data[:, 0]

    peak = np.abs(data).max()
    if peak > 1e-4:
        data = (data / peak * 0.95).astype(np.float32)

    out_path = os.path.join(tempfile.mkdtemp(), "output.wav")
    sf.write(out_path, data, sr)
    return out_path


if HAS_SPACES:
    @spaces.GPU(duration=120)
    def _generate_gpu(prompt, lyrics, audio_duration, infer_steps, seed):
        return _run_inference(prompt, lyrics, audio_duration, infer_steps, seed)
else:
    def _generate_gpu(prompt, lyrics, audio_duration, infer_steps, seed):
        return _run_inference(prompt, lyrics, audio_duration, infer_steps, seed)


# ── gr.Server App ─────────────────────────────────────────────────────────────
import gradio as gr
from gradio import Server
from fastapi.responses import HTMLResponse

app = Server(title="ace-studio")


@app.api(name="create", time_limit=300)
def create(description: str, audio_duration: float = 60.0, seed: int = -1) -> str:
    """One-box: describe a song → LLM composes tags+lyrics → our model generates audio.
    Returns JSON: {audio, title, tags, lyrics}"""
    try:
        composed = _compose(description)
        title, tags, lyrics = composed["title"], composed["tags"], composed["lyrics"]
        print(f"[create] title={title} tags={tags[:60]}...")

        wav_path = _generate_gpu(tags, lyrics, audio_duration, 8, seed)
        with open(wav_path, "rb") as f:
            wav_bytes = f.read()
        audio_b64 = f"data:audio/wav;base64,{base64.b64encode(wav_bytes).decode()}"

        return json.dumps({"audio": audio_b64, "title": title, "tags": tags, "lyrics": lyrics})
    except Exception as e:
        print(f"[create ERROR] {type(e).__name__}: {e}")
        print(traceback.format_exc())
        if "closed by visitor while queueing" in str(e).lower():
            raise RuntimeError(
                "The connection to the GPU queue was interrupted (this happens if the "
                "page reloads or the tab loses connection while waiting). Please try "
                "generating again without refreshing the page."
            ) from e
        raise


@app.api(name="generate", concurrency_limit=1, time_limit=180)
def generate(
    prompt: str,
    lyrics: str,
    audio_duration: float = 60.0,
    infer_step: int = 8,
    seed: int = -1,
) -> str:
    """Direct generate from explicit tags + lyrics (advanced mode). Returns base64 WAV data URL."""
    try:
        wav_path = _generate_gpu(prompt, lyrics, audio_duration, infer_step, seed)
        with open(wav_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode()
        return f"data:audio/wav;base64,{encoded}"
    except Exception as e:
        print(f"[generate ERROR] {type(e).__name__}: {e}")
        print(traceback.format_exc())
        raise


# ── Serve our custom HTML front page ──────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def homepage():
    with open(os.path.join(current_dir, "index.html"), "r", encoding="utf-8") as f:
        return f.read()


demo = app

if __name__ == "__main__":
    # مهم: حتماً demo.launch() (نه uvicorn دستی) — Server.launch() خودش یک Blocks
    # داخلی می‌سازد و blocks.launch() را صدا می‌زند، همان متدی که پکیج spaces پچ
    # می‌کند تا @spaces.GPU را در استارتاپ به ZeroGPU گزارش بدهد.
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        ssr_mode=False,
    )