"""
MusePlayer Backend Server
=========================
Self-hosted ACE-Step 1.5 music generation API on Hugging Face Spaces.
Deploy this as a Gradio Space. It serves both a Gradio UI and a REST API.

Requirements (requirements.txt):
------------------------------
torch>=2.5.0
diffusers>=0.38.0
transformers>=4.45.0
accelerate>=0.34.0
soundfile>=0.12.1
numpy>=1.26.0
gradio>=5.0.0
fastapi>=0.115.0
uvicorn>=0.30.0
pydantic>=2.9.0
requests>=2.32.0
pillow>=10.0.0
"""

import os
import io
import uuid
import json
import base64
import time
import threading
from datetime import datetime
from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field, asdict
from queue import Queue

import torch
import soundfile as sf
import numpy as np
from pydantic import BaseModel, Field
from fastapi import FastAPI, HTTPException, BackgroundTasks, Query
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import gradio as gr

# ------------------------------------------------------------------
#  Configuration
# ------------------------------------------------------------------
MODEL_ID = os.getenv("ACESTEP_MODEL", "ACE-Step/Ace-Step1.5")
DEVICE = os.getenv("ACESTEP_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
TORCH_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
MAX_QUEUE = int(os.getenv("MAX_QUEUE", 5))
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/tmp/museplayer_outputs")

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ------------------------------------------------------------------
#  Global state
# ------------------------------------------------------------------
pipe = None
generation_queue = Queue(maxsize=MAX_QUEUE)
results_store: Dict[str, Dict[str, Any]] = {}
user_prefs: Dict[str, Dict[str, Any]] = {}


# ------------------------------------------------------------------
#  Pydantic models
# ------------------------------------------------------------------
class GenerateRequest(BaseModel):
    prompt: str = Field(..., description="Music style description")
    lyrics: str = Field("", description="Lyrics with [verse]/[chorus] tags")
    duration: float = Field(30.0, ge=5.0, le=600.0, description="Duration in seconds")
    bpm: Optional[int] = Field(None, ge=40, le=300, description="Beats per minute")
    keyscale: Optional[str] = Field(None, description="e.g. C major, A minor")
    timesignature: Optional[str] = Field(None, description="e.g. 4 for 4/4")
    vocal_language: str = Field("en", description="Lyrics language code")
    num_inference_steps: int = Field(8, ge=1, le=100)
    guidance_scale: float = Field(7.0, ge=1.0, le=20.0)
    seed: Optional[int] = Field(None, description="Random seed for reproducibility")
    user_id: str = Field("anonymous", description="User identifier for personalization")
    station_type: str = Field("custom", description="Station genre/mood tag")


class FeedbackRequest(BaseModel):
    track_id: str
    user_id: str
    action: str = Field(..., pattern="^(like|dislike|skip|complete|favorite)$")
    station_type: Optional[str] = None


class UserProfile(BaseModel):
    user_id: str
    liked_prompts: List[str] = []
    disliked_prompts: List[str] = []
    favorite_genres: List[str] = []
    preferred_duration: float = 60.0
    preferred_bpm_range: List[int] = [80, 140]
    preferred_keys: List[str] = []


# ------------------------------------------------------------------
#  Model loading
# ------------------------------------------------------------------
def load_model():
    """Lazy-load the ACE-Step 1.5 pipeline."""
    global pipe
    if pipe is not None:
        return pipe

    from diffusers import AceStepPipeline

    print(f"[MusePlayer] Loading ACE-Step model: {MODEL_ID} on {DEVICE} ...")
    pipe = AceStepPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=TORCH_DTYPE,
    )
    pipe = pipe.to(DEVICE)
    if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_slicing"):
        pipe.vae.enable_slicing()
    if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"):
        pipe.vae.enable_tiling()
    print("[MusePlayer] Model loaded successfully.")
    return pipe


# ------------------------------------------------------------------
#  Music generation
# ------------------------------------------------------------------
def generate_music(
    prompt: str,
    lyrics: str = "",
    duration: float = 30.0,
    bpm: Optional[int] = None,
    keyscale: Optional[str] = None,
    timesignature: Optional[str] = None,
    vocal_language: str = "en",
    num_inference_steps: int = 8,
    guidance_scale: float = 7.0,
    seed: Optional[int] = None,
    track_id: Optional[str] = None,
) -> str:
    """Generate music and return the path to the output file."""
    model = load_model()

    generator = None
    if seed is not None:
        generator = torch.Generator(device=DEVICE).manual_seed(seed)

    kwargs = dict(
        prompt=prompt,
        lyrics=lyrics,
        audio_duration=duration,
        vocal_language=vocal_language,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        shift=3.0,
        generator=generator,
    )
    if bpm is not None:
        kwargs["bpm"] = bpm
    if keyscale is not None:
        kwargs["keyscale"] = keyscale
    if timesignature is not None:
        kwargs["timesignature"] = timesignature

    audio = model(**kwargs).audios

    # audio shape: [batch, channels, samples] or [batch, samples]
    # Save to file
    tid = track_id or str(uuid.uuid4())[:8]
    out_path = os.path.join(OUTPUT_DIR, f"{tid}.wav")
    audio_np = audio[0]
    if hasattr(audio_np, "cpu"):
        audio_np = audio_np.cpu()
    if hasattr(audio_np, "numpy"):
        audio_np = audio_np.numpy()
    if audio_np.ndim == 1:
        audio_np = audio_np[np.newaxis, :]
    sf.write(out_path, audio_np.T, model.sample_rate)
    return out_path


def worker_loop():
    """Background thread processing generation queue."""
    while True:
        try:
            job = generation_queue.get(timeout=1)
        except Exception:
            continue

        track_id = job["track_id"]
        try:
            start = time.time()
            path = generate_music(track_id=track_id, **job["params"])
            elapsed = time.time() - start

            results_store[track_id] = {
                "status": "completed",
                "track_id": track_id,
                "file_path": path,
                "created_at": datetime.utcnow().isoformat(),
                "generation_time_sec": round(elapsed, 2),
                "params": job["params"],
                "station_type": job.get("station_type", "custom"),
                "user_id": job.get("user_id", "anonymous"),
            }
        except Exception as e:
            results_store[track_id] = {
                "status": "failed",
                "track_id": track_id,
                "error": str(e),
                "params": job["params"],
            }
        finally:
            generation_queue.task_done()


# Start background worker
threading.Thread(target=worker_loop, daemon=True).start()


# ------------------------------------------------------------------
#  Preference / Personalization engine
# ------------------------------------------------------------------
PROMPT_TEMPLATES = {
    "lofi": "lo-fi hip hop, warm vinyl crackle, soft piano chords, dusty drums, relaxed, nostalgic, bedroom studio",
    "energetic": "upbeat electronic dance, driving four-on-the-floor kick, bright synth stabs, energetic, festival anthem",
    "chill": "ambient downtempo, airy pads, gentle acoustic guitar, warm bass, relaxed evening vibes, soft rain",
    "focus": "minimal instrumental, steady soft beat, warm synth pads, no vocals, productive focus, clean mix",
    "sleep": "slow ambient drone, deep soft pads, no percussion, gentle piano, sleep meditation, 432 Hz",
    "workout": "high-energy trap, aggressive 808s, fast hi-hats, powerful build-ups, gym motivation, intense",
    "acoustic": "intimate acoustic folk, fingerstyle guitar, warm vocals, gentle harmonica, campfire storytelling",
    "jazz": "smooth jazz, brushed drums, walking bass, muted trumpet, late night lounge, sophisticated",
    "classical": "cinematic orchestral, soaring strings, grand piano, emotional film score, epic crescendo",
    "synthwave": "retro synthwave, analog synth arpeggios, driving electronic beat, neon lights, 80s nostalgia",
}

BPM_RANGES = {
    "lofi": (60, 90),
    "energetic": (120, 140),
    "chill": (70, 100),
    "focus": (60, 90),
    "sleep": (40, 70),
    "workout": (130, 160),
    "acoustic": (80, 110),
    "jazz": (80, 120),
    "classical": (60, 120),
    "synthwave": (100, 130),
}


def get_user_profile(user_id: str) -> Dict[str, Any]:
    """Retrieve or initialize user preference profile."""
    if user_id not in user_prefs:
        user_prefs[user_id] = {
            "user_id": user_id,
            "liked_prompts": [],
            "disliked_prompts": [],
            "favorite_genres": [],
            "preferred_duration": 60.0,
            "preferred_bpm_range": [80, 140],
            "preferred_keys": [],
            "liked_tracks": [],
            "listening_history": [],
            "created_at": datetime.utcnow().isoformat(),
        }
    return user_prefs[user_id]


def update_preferences_from_feedback(user_id: str, action: str, track_meta: Dict[str, Any]):
    """Update user profile based on feedback."""
    profile = get_user_profile(user_id)
    prompt = track_meta.get("params", {}).get("prompt", "")
    station = track_meta.get("station_type", "custom")

    if action in ("like", "favorite"):
        if prompt and prompt not in profile["liked_prompts"]:
            profile["liked_prompts"].append(prompt)
        if station not in profile["favorite_genres"]:
            profile["favorite_genres"].append(station)
        profile["liked_tracks"].append(track_meta["track_id"])
    elif action == "dislike":
        if prompt and prompt not in profile["disliked_prompts"]:
            profile["disliked_prompts"].append(prompt)
    elif action == "complete":
        # Implicit positive signal
        pass

    # Keep history bounded
    profile["listening_history"].append({
        "track_id": track_meta.get("track_id"),
        "action": action,
        "station": station,
        "prompt": prompt,
        "timestamp": datetime.utcnow().isoformat(),
    })
    if len(profile["listening_history"]) > 200:
        profile["listening_history"] = profile["listening_history"][-200:]


def build_personalized_prompt(user_id: str, station_type: str = "custom") -> Dict[str, Any]:
    """Build a generation prompt tailored to the user's taste."""
    profile = get_user_profile(user_id)

    # Start from station template
    if station_type in PROMPT_TEMPLATES:
        base = PROMPT_TEMPLATES[station_type]
        bpm_low, bpm_high = BPM_RANGES[station_type]
    else:
        # Custom / blend from liked prompts
        if profile["liked_prompts"]:
            base = " ".join(profile["liked_prompts"][-3:])
            bpm_low, bpm_high = 80, 140
        else:
            base = "unique instrumental track, blended genre fusion, modern production"
            bpm_low, bpm_high = 80, 140

    # Inject positive style keywords from liked prompts
    liked_keywords = []
    for p in profile["liked_prompts"][-5:]:
        liked_keywords.extend(p.split(", ")[:3])

    if liked_keywords:
        liked_str = ", ".join(set(liked_keywords[-6:]))
        base = f"{base}. Also incorporate: {liked_str}"

    # Avoid disliked styles
    for dp in profile["disliked_prompts"][-3:]:
        if dp:
            base += f". Avoid: {dp.split(',')[0]} style"

    # Pick BPM in preferred range
    bpm = int((bpm_low + bpm_high) / 2)
    if profile["preferred_bpm_range"]:
        low, high = profile["preferred_bpm_range"]
        bpm = int(max(low, min(high, bpm)))

    duration = profile.get("preferred_duration", 60.0)

    return {
        "prompt": base,
        "duration": duration,
        "bpm": bpm,
        "station_type": station_type,
    }


# ------------------------------------------------------------------
#  FastAPI app
# ------------------------------------------------------------------
app = FastAPI(title="MusePlayer API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": pipe is not None, "queue_size": generation_queue.qsize()}


@app.get("/stations")
async def list_stations():
    """Return available station types with descriptions."""
    return {
        "stations": [
            {"id": k, "name": k.replace("_", " ").title(), "description": v[:80] + "..."}
            for k, v in PROMPT_TEMPLATES.items()
        ]
    }


@app.post("/generate")
async def api_generate(req: GenerateRequest, background_tasks: BackgroundTasks):
    """Queue a music generation job."""
    if generation_queue.full():
        raise HTTPException(status_code=503, detail="Generation queue is full. Try again shortly.")

    track_id = str(uuid.uuid4())[:12]
    params = {
        "prompt": req.prompt,
        "lyrics": req.lyrics,
        "duration": req.duration,
        "bpm": req.bpm,
        "keyscale": req.keyscale,
        "timesignature": req.timesignature,
        "vocal_language": req.vocal_language,
        "num_inference_steps": req.num_inference_steps,
        "guidance_scale": req.guidance_scale,
        "seed": req.seed,
    }

    job = {
        "track_id": track_id,
        "params": params,
        "user_id": req.user_id,
        "station_type": req.station_type,
    }

    # Pre-register result
    results_store[track_id] = {
        "status": "queued",
        "track_id": track_id,
        "queued_at": datetime.utcnow().isoformat(),
        "params": params,
    }
    generation_queue.put(job)

    return {"track_id": track_id, "status": "queued", "estimated_wait_sec": generation_queue.qsize() * 30}


@app.post("/generate-personalized")
async def api_generate_personalized(user_id: str = "anonymous", station_type: str = "lofi"):
    """Generate a track tailored to the user's taste profile."""
    if generation_queue.full():
        raise HTTPException(status_code=503, detail="Generation queue is full.")

    profile = get_user_profile(user_id)
    built = build_personalized_prompt(user_id, station_type)

    track_id = str(uuid.uuid4())[:12]
    params = {
        "prompt": built["prompt"],
        "lyrics": "",
        "duration": built["duration"],
        "bpm": built["bpm"],
        "keyscale": None,
        "timesignature": None,
        "vocal_language": "en",
        "num_inference_steps": 8,
        "guidance_scale": 7.0,
        "seed": None,
    }

    job = {
        "track_id": track_id,
        "params": params,
        "user_id": user_id,
        "station_type": station_type,
    }

    results_store[track_id] = {
        "status": "queued",
        "track_id": track_id,
        "queued_at": datetime.utcnow().isoformat(),
        "params": params,
        "personalized": True,
    }
    generation_queue.put(job)

    return {
        "track_id": track_id,
        "status": "queued",
        "prompt_used": built["prompt"],
        "estimated_wait_sec": generation_queue.qsize() * 30,
    }


@app.get("/track/{track_id}")
async def get_track(track_id: str):
    """Get track status or stream audio if complete."""
    if track_id not in results_store:
        raise HTTPException(status_code=404, detail="Track not found")

    meta = results_store[track_id]

    if meta["status"] == "completed":
        file_path = meta["file_path"]
        if not os.path.exists(file_path):
            raise HTTPException(status_code=404, detail="Audio file missing")

        def iterfile():
            with open(file_path, "rb") as f:
                yield from f

        return StreamingResponse(
            iterfile(),
            media_type="audio/wav",
            headers={"Content-Disposition": f'attachment; filename="{track_id}.wav"'},
        )

    return JSONResponse(content={"status": meta["status"], "track_id": track_id})


@app.get("/track/{track_id}/status")
async def track_status(track_id: str):
    if track_id not in results_store:
        raise HTTPException(status_code=404, detail="Track not found")
    meta = results_store[track_id]
    return {
        "status": meta["status"],
        "track_id": track_id,
        "params": meta.get("params"),
        "generation_time_sec": meta.get("generation_time_sec"),
        "error": meta.get("error"),
    }


@app.get("/track/{track_id}/download")
async def track_download(track_id: str, format: str = "wav"):
    """Download track in requested format (wav or mp3)."""
    if track_id not in results_store or results_store[track_id]["status"] != "completed":
        raise HTTPException(status_code=404, detail="Track not available")

    file_path = results_store[track_id]["file_path"]
    if format == "mp3":
        # Simple wav->mp3 conversion stub (requires ffmpeg in real deploy)
        mp3_path = file_path.replace(".wav", ".mp3")
        if not os.path.exists(mp3_path):
            # In production, use ffmpeg: os.system(f"ffmpeg -i {file_path} -q:a 2 {mp3_path}")
            # For HF Space, we'll return wav as fallback
            format = "wav"
            file_path = results_store[track_id]["file_path"]
        else:
            file_path = mp3_path

    def iterfile():
        with open(file_path, "rb") as f:
            yield from f

    return StreamingResponse(
        iterfile(),
        media_type="audio/wav" if format == "wav" else "audio/mpeg",
        headers={"Content-Disposition": f'attachment; filename="{track_id}.{format}"'},
    )


@app.post("/feedback")
async def api_feedback(req: FeedbackRequest):
    """Record user feedback for personalization."""
    track_id = req.track_id
    if track_id not in results_store:
        raise HTTPException(status_code=404, detail="Track not found")

    track_meta = results_store[track_id]
    update_preferences_from_feedback(req.user_id, req.action, track_meta)

    # Update station type if provided
    if req.station_type:
        track_meta["station_type"] = req.station_type

    return {"status": "ok", "user_id": req.user_id, "action": req.action}


@app.get("/profile/{user_id}")
async def get_profile(user_id: str):
    return get_user_profile(user_id)


@app.post("/profile/{user_id}")
async def update_profile(user_id: str, updates: Dict[str, Any]):
    profile = get_user_profile(user_id)
    allowed = {"preferred_duration", "preferred_bpm_range", "preferred_keys", "favorite_genres"}
    for k, v in updates.items():
        if k in allowed:
            profile[k] = v
    return profile


@app.get("/history/{user_id}")
async def get_history(user_id: str, limit: int = Query(20, ge=1, le=100)):
    profile = get_user_profile(user_id)
    history = profile.get("listening_history", [])
    return {"history": history[-limit:]}


@app.get("/tracks")
async def list_tracks():
    """List all generated tracks (for admin/demo)."""
    tracks = []
    for tid, meta in results_store.items():
        tracks.append({
            "track_id": tid,
            "status": meta["status"],
            "station_type": meta.get("station_type", "custom"),
            "user_id": meta.get("user_id", "anonymous"),
            "created_at": meta.get("created_at", meta.get("queued_at")),
            "prompt": meta.get("params", {}).get("prompt", "")[:60],
        })
    return {"tracks": sorted(tracks, key=lambda x: x.get("created_at", ""), reverse=True)}


# ------------------------------------------------------------------
#  Gradio UI (for HF Space demo + manual testing)
# ------------------------------------------------------------------
def gradio_generate(prompt, lyrics, duration, bpm, steps, guidance, seed):
    track_id = str(uuid.uuid4())[:8]
    out_path = generate_music(
        prompt=prompt,
        lyrics=lyrics,
        duration=duration,
        bpm=int(bpm) if bpm else None,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        seed=int(seed) if seed else None,
        track_id=track_id,
    )
    return out_path, f"Track ID: {track_id}"


def create_gradio_ui():
    with gr.Blocks(title="MusePlayer ACE-Step Server") as demo:
        gr.Markdown("# 🎵 MusePlayer Backend — ACE-Step 1.5")
        gr.Markdown("Self-hosted AI music generation. Use the REST API or test below.")

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    value="A beautiful piano piece with soft melodies and gentle rhythm",
                    lines=2,
                )
                lyrics = gr.Textbox(
                    label="Lyrics (optional)",
                    placeholder="[verse]\nSoft notes in the morning light\n[chorus]\nMusic fills the air tonight",
                    lines=4,
                )
                duration = gr.Slider(5, 120, value=30, step=5, label="Duration (seconds)")
                bpm = gr.Number(value=120, label="BPM (optional)")
                steps = gr.Slider(1, 50, value=8, step=1, label="Inference Steps")
                guidance = gr.Slider(1.0, 20.0, value=7.0, step=0.5, label="Guidance Scale")
                seed = gr.Number(value=None, label="Seed (optional)")
                btn = gr.Button("Generate Music", variant="primary")

            with gr.Column():
                audio_out = gr.Audio(label="Generated Music", type="filepath")
                info_out = gr.Textbox(label="Info", interactive=False)

        btn.click(
            gradio_generate,
            inputs=[prompt, lyrics, duration, bpm, steps, guidance, seed],
            outputs=[audio_out, info_out],
        )

        gr.Markdown("""
        ### API Endpoints
        - `POST /generate` — Queue a track
        - `POST /generate-personalized` — Auto-tailored track
        - `GET /track/{id}` — Stream audio
        - `GET /track/{id}/status` — Check status
        - `POST /feedback` — Like / dislike / skip
        - `GET /profile/{user_id}` — Get user taste profile
        - `GET /stations` — List stations
        """)
    return demo


# Mount Gradio into FastAPI so both run on the same port
demo = create_gradio_ui()
app = gr.mount_gradio_app(app, demo, path="/")


# ------------------------------------------------------------------
#  Entry point
# ------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    # Preload model on startup if env says so
    if os.getenv("PRELOAD_MODEL", "1") == "1":
        threading.Thread(target=load_model, daemon=True).start()
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)