import io
import time
import queue
import threading
import torch
import scipy.io.wavfile as wavfile
import uvicorn
from transformers import AutoProcessor, MusicgenForConditionalGeneration, LogitsProcessor, LogitsProcessorList
import gradio as gr
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

print("[musicgen] 🚀 App starting...")

torch.set_num_threads(2)
torch.set_num_interop_threads(1)

_model = None
_processor = None
_lock = threading.Lock()
_loaded = False
_error = None


def _load():
    global _model, _processor, _loaded, _error
    if _loaded:
        return
    try:
        print("[musicgen] ⏳ Loading facebook/musicgen-medium ...")
        t0 = time.time()
        _processor = AutoProcessor.from_pretrained("facebook/musicgen-medium")
        _model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-medium")
        _model.eval()
        _loaded = True
        print(f"[musicgen] ✅ Model loaded in {time.time()-t0:.1f}s")
    except Exception as e:
        _error = str(e)
        print(f"[musicgen] ❌ Load failed: {e}")
        raise


threading.Thread(target=_load, daemon=True).start()


# ── 进度追踪：每生成 10 个 token 回调一次 ──────────────────
class ProgressTracker(LogitsProcessor):
    def __init__(self, max_tokens: int, log_fn):
        self.max_tokens = max_tokens
        self.generated = 0
        self.log_fn = log_fn
        self.start_time = time.time()

    def __call__(self, input_ids, scores):
        self.generated += 1
        if self.generated % 10 == 0 or self.generated == self.max_tokens:
            elapsed = time.time() - self.start_time
            pct = self.generated / self.max_tokens * 100
            eta = (elapsed / self.generated) * (self.max_tokens - self.generated) if self.generated > 0 else 0
            self.log_fn(
                f"[{pct:5.1f}%] token {self.generated:>4}/{self.max_tokens} | "
                f"elapsed {elapsed:>5.0f}s | ETA ~{eta:.0f}s"
            )
        return scores


def _generate(prompt: str, duration: int = 8, guidance: float = 3.0, log_fn=None):
    if not _loaded:
        _load()
    max_new_tokens = int(duration * 50)
    processors = LogitsProcessorList([ProgressTracker(max_new_tokens, log_fn)]) if log_fn else None
    with _lock:
        inputs = _processor(text=[prompt], padding=True, return_tensors="pt")
        with torch.no_grad():
            audio_values = _model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                guidance_scale=guidance,
                logits_processor=processors,
            )
    sr = _model.config.audio_encoder.sampling_rate
    audio_np = audio_values[0, 0].numpy()
    return sr, audio_np


# ── Gradio UI ────────────────────────────────────────
def get_status():
    if _error:
        return f"**Status:** ❌ Load failed: {_error}"
    if _loaded:
        return "**Status:** ✅ Model ready — you can generate now!"
    return "**Status:** ⏳ Model loading, please wait..."


def ui_generate(prompt, duration, guidance):
    if not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if not _loaded:
        raise gr.Error("Model is still loading. Please retry.")

    log_q = queue.Queue()
    result = {}

    def log_fn(msg):
        log_q.put(msg)

    def run():
        try:
            sr, audio_np = _generate(prompt, int(duration), float(guidance), log_fn=log_fn)
            result["sr"] = sr
            result["audio"] = audio_np
        except Exception as e:
            result["error"] = str(e)
        finally:
            log_q.put(None)  # 结束信号

    threading.Thread(target=run, daemon=True).start()

    t0 = time.time()
    max_tokens = int(duration * 50)
    log_lines = [f"[  0.0%] Starting generation — {max_tokens} tokens total ({duration}s audio)..."]
    yield None, "⏳ Generating...", "\n".join(log_lines)

    while True:
        try:
            msg = log_q.get(timeout=30)
        except queue.Empty:
            log_lines.append("[WARN] No response for 30s, still waiting...")
            yield None, "⏳ Generating...", "\n".join(log_lines)
            continue
        if msg is None:
            break
        log_lines.append(msg)
        yield None, "⏳ Generating...", "\n".join(log_lines)

    if "error" in result:
        raise gr.Error(result["error"])

    elapsed = time.time() - t0
    log_lines.append(f"[100.0%] ✅ Done in {elapsed:.1f}s")
    yield (result["sr"], result["audio"]), f"✅ Done in {elapsed:.1f}s", "\n".join(log_lines)


with gr.Blocks(title="MusicGen") as demo:
    gr.Markdown("# 🎵 MusicGen — Text to Music")
    status = gr.Markdown(get_status())
    prompt = gr.Textbox(label="Prompt", placeholder="upbeat pop song with electric guitar")
    with gr.Row():
        dur = gr.Slider(2, 30, value=8, step=1, label="Duration (s)")
        guide = gr.Slider(1.0, 5.0, value=3.0, step=0.1, label="Guidance")
    btn = gr.Button("🎵 Generate", variant="primary")
    audio = gr.Audio(label="Result", type="numpy")
    msg = gr.Markdown("")
    log_box = gr.Textbox(label="📋 Generation Log", lines=12, interactive=False, max_lines=20)
    btn.click(ui_generate, [prompt, dur, guide], [audio, msg, log_box])

    timer = gr.Timer(5)
    timer.tick(get_status, outputs=status)


# ── 关键修复：自己建 FastAPI app，把自定义路由和 Gradio 挂到同一个 app 上 ──
# 之前直接用 demo.app 注册路由 + demo.launch() 起服务，
# 两者不是同一个 FastAPI 实例，导致 /health /generate 全部 404。
# 现在改为：自建 app → 注册自定义路由 → mount_gradio_app 挂 Gradio UI →
# uvicorn.run(app) 统一对外提供服务。

app = FastAPI()


class GenerateRequest(BaseModel):
    prompt: str
    duration: int = 8
    guidance: float = 3.0


@app.get("/health")
def health():
    return {"status": "ready" if _loaded else "loading", "error": _error}


@app.post("/generate")
def generate(req: GenerateRequest):
    if not req.prompt.strip():
        raise HTTPException(status_code=400, detail="prompt is required")
    if not _loaded:
        raise HTTPException(status_code=503, detail="Model still loading, retry later")
    try:
        sr, audio_np = _generate(req.prompt, req.duration, req.guidance)
        buf = io.BytesIO()
        wavfile.write(buf, sr, audio_np)
        buf.seek(0)
        return StreamingResponse(buf, media_type="audio/wav",
                                 headers={"Content-Disposition": "attachment; filename=output.wav"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# CORS：允许你的网站域名跨域调用 /generate /health
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # 如需限制改成你的网站域名，如 ["https://yoursite.com"]
    allow_methods=["GET", "POST"],
    allow_headers=["Content-Type"],
)

# 把 Gradio UI 挂载到根路径 "/"，和上面的自定义路由共用同一个 app 实例
app = gr.mount_gradio_app(app, demo, path="/")


if __name__ == "__main__":
    # 不再用 demo.launch()，改为直接用 uvicorn 运行合并后的 app，
    # 确保 /health /generate 和 Gradio UI 是同一个进程里的同一个 app。
    uvicorn.run(app, host="0.0.0.0", port=7860)