"""Bytley Demucs stem-separation backend (Hugging Face Space, free CPU tier).

POST /separate  (multipart: file; query: mode=stems|karaoke)
  → application/zip of MP3 stems:
    mode=stems   → vocals.mp3, drums.mp3, bass.mp3, other.mp3
    mode=karaoke → instrumental.mp3, vocals.mp3

Uses the Demucs CLI (stable across versions) with the htdemucs model baked into the
image. Files are processed in a temp dir and deleted right after the response.
"""
import io
import os
import subprocess
import sys
import tempfile
import zipfile

from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response

MAX_MB = int(os.environ.get("MAX_MB", "30"))
TOKEN = os.environ.get("BYTLEY_TOKEN", "")

app = FastAPI(title="bytley-demucs")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://bytley.co",
        "https://www.bytley.co",
        "http://localhost:4321",
        "http://localhost:4322",
    ],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


@app.get("/")
def health():
    return {"ok": True, "service": "demucs-stems", "model": "htdemucs", "maxMb": MAX_MB}


def _run(cmd: list[str], timeout: int):
    r = subprocess.run(cmd, capture_output=True, timeout=timeout)
    if r.returncode != 0:
        stderr_tail = r.stderr.decode(errors="replace")[-800:]
        raise HTTPException(
            500,
            detail=f"Separation failed. stderr tail: {stderr_tail}",
        )


@app.post("/separate")
async def separate(
    file: UploadFile = File(...),
    mode: str = "stems",
    x_bytley_token: str | None = Header(default=None),
):
    if TOKEN and x_bytley_token != TOKEN:
        raise HTTPException(401, detail="Unauthorized.")
    data = await file.read()
    if len(data) > MAX_MB * 1024 * 1024:
        raise HTTPException(413, detail=f"File too large — max {MAX_MB} MB.")
    if len(data) < 1024:
        raise HTTPException(400, detail="Empty file.")

    with tempfile.TemporaryDirectory() as td:
        ext = os.path.splitext(file.filename or "a.mp3")[1] or ".mp3"
        src = os.path.join(td, "input" + ext)
        wav = os.path.join(td, "track.wav")
        with open(src, "wb") as f:
            f.write(data)

        # Normalize to 44.1 kHz stereo WAV (also validates the upload).
        r = subprocess.run(
            ["ffmpeg", "-y", "-i", src, "-vn", "-ac", "2", "-ar", "44100", "-t", "600", "-c:a", "pcm_s16le", wav],
            capture_output=True,
            timeout=300,
        )
        if r.returncode != 0 or not os.path.exists(wav):
            raise HTTPException(400, detail="Couldn't read audio from this file.")

        # Demucs CLI — stable across package versions.
        outdir = os.path.join(td, "sep")
        # demucs v4 entry point is `python -m demucs`, NOT `python -m demucs.separate`
        cmd = [sys.executable, "-m", "demucs", "-n", "htdemucs", "-d", "cpu", "-o", outdir]
        if mode == "karaoke":
            cmd += ["--two-stems", "vocals"]
        cmd.append(wav)
        _run(cmd, timeout=1700)

        stemdir = os.path.join(outdir, "htdemucs", "track")
        if not os.path.isdir(stemdir):
            raise HTTPException(500, detail="Separation produced no output.")

        names = ["instrumental", "vocals"] if mode == "karaoke" else ["vocals", "drums", "bass", "other"]
        src_names = {"instrumental": "no_vocals"}  # demucs names the karaoke stem no_vocals

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as z:
            for name in names:
                wpath = os.path.join(stemdir, f"{src_names.get(name, name)}.wav")
                if not os.path.exists(wpath):
                    raise HTTPException(500, detail=f"Missing stem: {name}.")
                mpath = os.path.join(td, f"{name}.mp3")
                subprocess.run(
                    ["ffmpeg", "-y", "-i", wpath, "-c:a", "libmp3lame", "-b:a", "160k", mpath],
                    capture_output=True,
                    timeout=300,
                    check=True,
                )
                z.write(mpath, arcname=f"{name}.mp3")

    return Response(
        content=buf.getvalue(),
        media_type="application/zip",
        headers={"Content-Disposition": 'attachment; filename="stems.zip"'},
    )