app.py

import json
import os
import shutil
import subprocess
import tempfile
import threading
import time
import uuid
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request as UrlRequest, urlopen, urlretrieve
from zipfile import ZIP_DEFLATED, ZipFile

from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from starlette.background import BackgroundTask


APP_NAME = "Neuralis Stem Worker"
DEFAULT_MODEL = os.getenv("NEURALIS_DEMUCS_MODEL", "htdemucs")
DEFAULT_MODE = os.getenv("NEURALIS_STEM_MODE", "fast-2stem")
DEFAULT_FORMAT = os.getenv("NEURALIS_STEM_FORMAT", "mp3")
MAX_UPLOAD_MB = int(os.getenv("NEURALIS_STEM_MAX_UPLOAD_MB", "300"))
MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024
ALLOWED_MODELS = {
    "htdemucs",
    "htdemucs_ft",
    "htdemucs_6s",
    "uvr_mdx_voc_ft",
}
UVR_MODEL_FILES = {
    "uvr_mdx_voc_ft": "UVR-MDX-NET-Voc_FT.onnx",
}
MODEL_LABELS = {
    "htdemucs": "Demucs Standard",
    "htdemucs_ft": "Demucs Fine-Tuned",
    "htdemucs_6s": "Demucs 6 Stem",
    "uvr_mdx_voc_ft": "UVR MDX Vocal FT",
}
ALLOWED_MODES = {
    "fast-2stem",
    "premium-4stem",
}
ALLOWED_FORMATS = {
    "wav",
    "mp3",
}
JOB_TTL_SECONDS = 2 * 60 * 60
DEFAULT_UVR_WORKER_URL = "https://jayman-neuralis-uvr-stem-worker.hf.space"
DEFAULT_TEST_API_KEY = "neuralis-stem-test-2026"

app = FastAPI(title=APP_NAME)
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.getenv("NEURALIS_STEM_CORS_ORIGINS", "*").split(","),
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)
JOBS = {}
JOBS_LOCK = threading.Lock()


def _client_key(request: Request) -> str:
    auth = request.headers.get("authorization", "")
    if auth.lower().startswith("bearer "):
        return auth[7:].strip()
    return request.headers.get("x-neuralis-api-key", "").strip()


def _require_api_key(request: Request) -> None:
    expected = os.getenv("NEURALIS_STEM_API_KEY", "").strip()
    if not expected:
        return
    if _client_key(request) != expected:
        raise HTTPException(status_code=401, detail="Invalid Neuralis stem API key")


def _check_api_key(request: Request, form_key: str = "") -> None:
    expected = os.getenv("NEURALIS_STEM_API_KEY", "").strip()
    if not expected:
        return
    header_key = _client_key(request)
    if header_key and header_key != expected:
        raise HTTPException(status_code=401, detail="Invalid Neuralis stem API key")
    if not header_key and form_key.strip() != expected:
        raise HTTPException(status_code=401, detail="Invalid Neuralis stem API key")


def _set_job(job_id: str, **updates) -> None:
    with JOBS_LOCK:
        job = JOBS.get(job_id)
        if not job:
            return
        job.update(updates)
        job["updatedAt"] = time.time()


def _public_job(job: dict) -> dict:
    status = job.get("status", "queued")
    progress = float(job.get("progress", 0))
    if status == "processing":
        elapsed = max(0.0, time.time() - float(job.get("startedAt", time.time())))
        estimate = 120.0 if job.get("mode") == "fast-2stem" else 240.0
        progress = max(progress, min(92.0, 18.0 + (elapsed / estimate) * 70.0))
    return {
        "id": job["id"],
        "status": status,
        "progress": round(progress, 1),
        "stage": job.get("stage", ""),
        "mode": job.get("mode", DEFAULT_MODE),
        "model": job.get("model", DEFAULT_MODEL),
        "format": job.get("format", DEFAULT_FORMAT),
        "source": job.get("source", ""),
        "downloadUrl": job.get("downloadUrl"),
        "error": job.get("error"),
    }


def _cleanup_old_jobs() -> None:
    cutoff = time.time() - JOB_TTL_SECONDS
    stale = []
    with JOBS_LOCK:
        for job_id, job in JOBS.items():
            if float(job.get("createdAt", 0)) < cutoff:
                stale.append((job_id, job.get("workDir")))
        for job_id, _ in stale:
            JOBS.pop(job_id, None)
    for _, work_dir in stale:
        if work_dir:
            shutil.rmtree(work_dir, ignore_errors=True)


def _safe_name(filename: str) -> str:
    name = Path(filename or "upload.wav").name
    keep = []
    for ch in name:
        if ch.isalnum() or ch in (" ", ".", "-", "_"):
            keep.append(ch)
    cleaned = "".join(keep).strip(" .")
    return cleaned or "upload.wav"


async def _save_upload(upload: UploadFile, target: Path) -> int:
    total = 0
    with target.open("wb") as out:
        while True:
            chunk = await upload.read(1024 * 1024)
            if not chunk:
                break
            total += len(chunk)
            if total > MAX_UPLOAD_BYTES:
                raise HTTPException(
                    status_code=413,
                    detail=f"Upload is larger than {MAX_UPLOAD_MB} MB",
                )
            out.write(chunk)
    return total


def _run_demucs(input_path: Path, work_dir: Path, model: str, mode: str) -> dict:
    output_dir = work_dir / "separated"
    cmd = [
        "python",
        "-m",
        "demucs.separate",
        "--name",
        model,
        "--out",
        str(output_dir),
        "--filename",
        "{stem}.{ext}",
    ]
    if mode == "fast-2stem":
        cmd.extend(["--two-stems", "vocals"])
    cmd.append(str(input_path))

    start = time.time()
    proc = subprocess.run(
        cmd,
        cwd=work_dir,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        timeout=60 * 20,
    )
    if proc.returncode != 0:
        raise RuntimeError(proc.stdout[-6000:])

    stem_dir = output_dir / model
    if not stem_dir.exists():
        candidates = [p for p in output_dir.rglob("*") if p.is_dir()]
        if candidates:
            stem_dir = candidates[-1]

    if mode == "fast-2stem":
        required = ["vocals.wav", "no_vocals.wav"]
        output_stems = [
            ("vocals.wav", "vocals.wav"),
            ("no_vocals.wav", "instrumental.wav"),
        ]
    else:
        required = ["vocals.wav", "drums.wav", "bass.wav", "other.wav"]
        output_stems = [(name, name) for name in required]

    missing = [name for name in required if not (stem_dir / name).exists()]
    if missing:
        raise RuntimeError(f"Demucs finished but stems are missing: {', '.join(missing)}")

    elapsed = time.time() - start
    (work_dir / "neuralis-stem-report.txt").write_text(
        f"mode={mode}\nmodel={model}\nseconds={elapsed:.2f}\nsource={input_path.name}\n",
        encoding="utf-8",
    )
    return {
        "stemDir": stem_dir,
        "outputStems": output_stems,
        "elapsed": elapsed,
    }


def _score_uvr_stem(path: Path, kind: str) -> int:
    name = path.name.lower()
    score = 0
    if kind == "vocal":
        if "vocals" in name:
            score += 12
        if "vocal" in name:
            score += 8
        if "instrumental" in name or "inst" in name or "no_vocal" in name:
            score -= 20
    else:
        if "instrumental" in name:
            score += 12
        if "inst" in name:
            score += 8
        if "no_vocal" in name or "novocal" in name:
            score += 10
        if "vocals" in name and "no_vocal" not in name and "novocal" not in name:
            score -= 20
    if "converted" in name:
        score -= 2
    return score


def _find_uvr_stems(output_dir: Path) -> tuple[Path, Path]:
    candidates = [p for p in output_dir.rglob("*") if p.is_file() and p.suffix.lower() in {".wav", ".mp3", ".flac"}]
    if len(candidates) < 2:
        raise RuntimeError("UVR MDX finished but did not produce two stems")
    vocal = max(candidates, key=lambda p: _score_uvr_stem(p, "vocal"))
    instrumental_pool = [p for p in candidates if p != vocal]
    instrumental = max(instrumental_pool, key=lambda p: _score_uvr_stem(p, "instrumental"))
    if _score_uvr_stem(vocal, "vocal") <= 0 or _score_uvr_stem(instrumental, "instrumental") <= 0:
        names = ", ".join(p.name for p in candidates[:8])
        raise RuntimeError(f"UVR MDX stems could not be identified from outputs: {names}")
    return vocal, instrumental


def _normalize_uvr_worker_url(value: str) -> str:
    raw = (value or DEFAULT_UVR_WORKER_URL).strip()
    marker = "huggingface.co/spaces/"
    if marker in raw:
        repo = raw.split(marker, 1)[1].split("?", 1)[0].split("#", 1)[0].strip("/")
        parts = repo.split("/")
        if len(parts) >= 2:
            return f"https://{parts[0]}-{parts[1]}.hf.space"
    return raw or DEFAULT_UVR_WORKER_URL


def _flatten_gradio_outputs(value) -> list[str]:
    paths = []
    if value is None:
        return paths
    if isinstance(value, (str, Path)):
        return [str(value)]
    if isinstance(value, dict):
        for key in ("path", "name", "url"):
            item = value.get(key)
            if item:
                paths.append(str(item))
        for item in value.values():
            if isinstance(item, (dict, list, tuple)):
                paths.extend(_flatten_gradio_outputs(item))
        return paths
    if isinstance(value, (list, tuple)):
        for item in value:
            paths.extend(_flatten_gradio_outputs(item))
    return paths


def _copy_gradio_output(value: str, output_dir: Path, index: int) -> Path | None:
    raw = str(value or "").strip()
    if not raw:
        return None
    suffix = Path(raw.split("?", 1)[0]).suffix.lower()
    if suffix not in {".wav", ".mp3", ".flac"}:
        suffix = ".wav"
    raw_name = Path(raw.split("?", 1)[0]).name
    target_name = _safe_name(raw_name or f"uvr-output-{index}{suffix}")
    if Path(target_name).suffix.lower() not in {".wav", ".mp3", ".flac"}:
        target_name = f"{Path(target_name).stem}{suffix}"
    target = output_dir / target_name
    if target.exists():
        target = output_dir / f"{Path(target_name).stem}-{index}{Path(target_name).suffix}"
    if raw.startswith(("http://", "https://")):
        urlretrieve(raw, target)
        return target
    source = Path(raw)
    if source.exists() and source.is_file():
        if source.resolve() == target.resolve():
            return target
        shutil.copy2(source, target)
        return target
    return None


def _prepare_ordered_uvr_stems(copied: list[Path], output_dir: Path) -> tuple[Path, Path]:
    audio_files = [p for p in copied if p.exists() and p.suffix.lower() in {".wav", ".mp3", ".flac"}]
    if len(audio_files) >= 2:
        best_vocal = max(audio_files, key=lambda p: _score_uvr_stem(p, "vocal"))
        best_music = max([p for p in audio_files if p != best_vocal], key=lambda p: _score_uvr_stem(p, "instrumental"))
        if _score_uvr_stem(best_vocal, "vocal") > 0 and _score_uvr_stem(best_music, "instrumental") > 0:
            vocal_source, music_source = best_vocal, best_music
        else:
            vocal_source, music_source = audio_files[0], audio_files[1]
        vocal_target = output_dir / f"vocals{vocal_source.suffix.lower()}"
        music_target = output_dir / f"instrumental{music_source.suffix.lower()}"
        if vocal_source.resolve() != vocal_target.resolve():
            shutil.copy2(vocal_source, vocal_target)
        if music_source.resolve() != music_target.resolve():
            shutil.copy2(music_source, music_target)
        return vocal_target, music_target
    return _find_uvr_stems(output_dir)


def _uvr_url(worker_url: str, path: str) -> str:
    return f"{worker_url.rstrip('/')}/{path.lstrip('/')}"


def _uvr_api_key() -> str:
    return (
        os.getenv("NEURALIS_UVR_WORKER_API_KEY", "").strip()
        or os.getenv("NEURALIS_STEM_API_KEY", "").strip()
        or DEFAULT_TEST_API_KEY
    )


def _read_json_url(url: str, timeout: int = 30) -> dict:
    with urlopen(url, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8", errors="replace"))


def _post_uvr_multipart(worker_url: str, path: str, input_path: Path, model: str, mode: str) -> bytes:
    boundary = f"----neuralis-uvr-{uuid.uuid4().hex}"
    api_key = _uvr_api_key()
    fields = {
        "model": "uvr_mdx_voc_ft",
        "mode": "fast-2stem",
        "format": "wav",
    }
    if api_key:
        fields["apiKey"] = api_key
    parts = []
    for name, value in fields.items():
        parts.append(
            f"--{boundary}\r\n"
            f'Content-Disposition: form-data; name="{name}"\r\n\r\n'
            f"{value}\r\n".encode("utf-8")
        )
    filename = _safe_name(input_path.name)
    parts.append(
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
        "Content-Type: application/octet-stream\r\n\r\n".encode("utf-8")
        + input_path.read_bytes()
        + b"\r\n"
    )
    parts.append(f"--{boundary}--\r\n".encode("utf-8"))
    headers = {
        "Content-Type": f"multipart/form-data; boundary={boundary}",
        "Accept": "application/json",
    }
    if api_key:
        headers["X-Neuralis-API-Key"] = api_key
    request = UrlRequest(
        _uvr_url(worker_url, path),
        data=b"".join(parts),
        headers=headers,
        method="POST",
    )
    try:
        with urlopen(request, timeout=120) as response:
            return response.read()
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:900]
        raise RuntimeError(f"UVR worker {path} returned HTTP {exc.code}: {detail}") from exc


def _extract_uvr_job_id(value) -> str:
    if not isinstance(value, dict):
        return ""
    for key in ("id", "jobId", "job_id"):
        if value.get(key):
            return str(value.get(key))
    nested = value.get("job")
    if isinstance(nested, dict):
        return _extract_uvr_job_id(nested)
    return ""


def _post_uvr_job_payload(worker_url: str, input_path: Path, model: str, mode: str) -> bytes:
    return _post_uvr_multipart(worker_url, "/jobs", input_path, model, mode)


def _post_uvr_job(worker_url: str, input_path: Path, model: str, mode: str) -> dict:
    return json.loads(_post_uvr_job_payload(worker_url, input_path, model, mode).decode("utf-8", errors="replace"))


def _prepare_uvr_zip_payload(payload: bytes, work_dir: Path, name: str = "uvr-worker-stems.zip") -> dict:
    zip_path = work_dir / "uvr-worker-direct-stems.zip"
    if name:
        zip_path = work_dir / name
    zip_path.write_bytes(payload)
    output_dir = work_dir / "uvr-worker-direct-separated"
    output_dir.mkdir(parents=True, exist_ok=True)
    with ZipFile(zip_path, "r") as archive:
        archive.extractall(output_dir)
    vocal_path, instrumental_path = _find_uvr_stems(output_dir)
    return {
        "stemDir": output_dir,
        "outputStems": [
            (vocal_path.relative_to(output_dir).as_posix(), "vocals.wav"),
            (instrumental_path.relative_to(output_dir).as_posix(), "instrumental.wav"),
        ],
        "elapsed": 0.0,
    }


def _try_uvr_separate(worker_url: str, input_path: Path, work_dir: Path, model: str, mode: str) -> dict | None:
    try:
        payload = _post_uvr_multipart(worker_url, "/separate", input_path, model, mode)
    except RuntimeError as exc:
        if "HTTP 404" in str(exc) or "HTTP 405" in str(exc):
            return None
        raise
    return _prepare_uvr_zip_payload(payload, work_dir)


def _try_neuralis_uvr_worker(input_path: Path, work_dir: Path, model: str, mode: str, worker_url: str) -> dict | None:
    try:
        payload = _post_uvr_job_payload(worker_url, input_path, model, mode)
        try:
            job = json.loads(payload.decode("utf-8", errors="replace"))
        except json.JSONDecodeError:
            return _prepare_uvr_zip_payload(payload, work_dir, "uvr-worker-jobs-response.zip")
        job_id = _extract_uvr_job_id(job)
        if not job_id:
            raise RuntimeError(f"UVR worker /jobs returned no id: {json.dumps(job, ensure_ascii=True)[:700]}")
        deadline = time.time() + 60 * 30
        status = job
        while time.time() < deadline:
            time.sleep(2.0)
            status = _read_json_url(_uvr_url(worker_url, f"/jobs/{job_id}"), timeout=30)
            if status.get("status") == "ready":
                break
            if status.get("status") == "failed":
                raise RuntimeError(status.get("error") or "UVR worker failed")
        if status.get("status") != "ready":
            raise RuntimeError("UVR worker timed out")

        zip_url = status.get("downloadUrl") or f"/jobs/{job_id}/download"
        if str(zip_url).startswith(("http://", "https://")):
            download_url = str(zip_url)
        else:
            download_url = _uvr_url(worker_url, zip_url)
        zip_path = work_dir / "uvr-worker-stems.zip"
        urlretrieve(download_url, zip_path)

        output_dir = work_dir / "uvr-worker-separated"
        output_dir.mkdir(parents=True, exist_ok=True)
        with ZipFile(zip_path, "r") as archive:
            archive.extractall(output_dir)
        vocal_path, instrumental_path = _find_uvr_stems(output_dir)
        elapsed = max(0.0, time.time() - float(status.get("createdAt", 0))) if status.get("createdAt") else 0.0
        (work_dir / "neuralis-stem-report.txt").write_text(
            f"mode={mode}\nmodel={model}\nengine=remote-neuralis-uvr\nworker={worker_url}\nseconds={elapsed:.2f}\nsource={input_path.name}\n",
            encoding="utf-8",
        )
        return {
            "stemDir": output_dir,
            "outputStems": [
                (vocal_path.relative_to(output_dir).as_posix(), "vocals.wav"),
                (instrumental_path.relative_to(output_dir).as_posix(), "instrumental.wav"),
            ],
            "elapsed": elapsed,
        }
    except RuntimeError as exc:
        if "HTTP 404" in str(exc) or "HTTP 405" in str(exc):
            return None
        raise
    except (URLError, TimeoutError, json.JSONDecodeError) as exc:
        raise RuntimeError(f"Could not reach the UVR worker API: {exc}") from exc


def _run_audio_separator(input_path: Path, work_dir: Path, model: str, mode: str) -> dict:
    if mode != "fast-2stem":
        raise RuntimeError("UVR MDX test model currently supports fast-2stem vocal/instrumental output only")
    model_file = UVR_MODEL_FILES.get(model)
    if not model_file:
        raise RuntimeError(f"Unsupported UVR model: {model}")

    worker_url = _normalize_uvr_worker_url(os.getenv("NEURALIS_UVR_WORKER_URL", DEFAULT_UVR_WORKER_URL))
    try:
        neuralis_result = _try_neuralis_uvr_worker(input_path, work_dir, model, mode, worker_url)
        if neuralis_result:
            return neuralis_result
    except Exception as job_error:
        try:
            direct_result = _try_uvr_separate(worker_url, input_path, work_dir, model, mode)
            if direct_result:
                return direct_result
        except Exception as direct_error:
            raise RuntimeError(f"UVR worker request failed. jobs: {job_error}; separate: {direct_error}") from direct_error
        raise

    direct_result = _try_uvr_separate(worker_url, input_path, work_dir, model, mode)
    if direct_result:
        return direct_result
    raise RuntimeError("UVR worker returned no job id and no direct stem package")


def _run_separator(input_path: Path, work_dir: Path, model: str, mode: str) -> dict:
    if model in UVR_MODEL_FILES:
        return _run_audio_separator(input_path, work_dir, model, mode)
    return _run_demucs(input_path, work_dir, model, mode)


def _convert_stem_to_mp3(source_path: Path, target_path: Path) -> None:
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(source_path),
        "-codec:a",
        "libmp3lame",
        "-b:a",
        "320k",
        str(target_path),
    ]
    proc = subprocess.run(
        cmd,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        timeout=60 * 5,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"MP3 conversion failed: {proc.stdout[-2000:]}")


def _make_zip(run_info: dict, work_dir: Path, original_name: str, model: str, mode: str, output_format: str) -> Path:
    stem_dir = run_info["stemDir"]
    output_stems = run_info["outputStems"]
    zip_path = work_dir / "neuralis-stems.zip"
    archive_stems = []
    with ZipFile(zip_path, "w", ZIP_DEFLATED) as archive:
        for source_name, archive_name in output_stems:
            source_path = stem_dir / source_name
            if output_format == "mp3":
                archive_name = f"{Path(archive_name).stem}.mp3"
                mp3_path = work_dir / archive_name
                _convert_stem_to_mp3(source_path, mp3_path)
                archive.write(mp3_path, archive_name)
            else:
                archive.write(source_path, archive_name)
            archive_stems.append(archive_name)
        report = work_dir / "neuralis-stem-report.txt"
        if report.exists():
            archive.write(report, "neuralis-stem-report.txt")
        archive.writestr(
            "manifest.json",
            json.dumps(
                {
                    "source": original_name,
                    "mode": mode,
                    "model": model,
                    "format": output_format,
                    "stems": archive_stems,
                    "seconds": round(float(run_info["elapsed"]), 2),
                },
                indent=2,
            ),
        )
    return zip_path


def _process_job(job_id: str, input_path: Path, work_dir: Path, original_name: str, model: str, mode: str, output_format: str) -> None:
    try:
        _set_job(
            job_id,
            status="processing",
            progress=16,
            stage=f"Loading {MODEL_LABELS.get(model, model)} separation model",
            startedAt=time.time(),
        )
        run_info = _run_separator(input_path, work_dir, model, mode)
        pack_stage = "Encoding MP3 stems" if output_format == "mp3" else "Packing stems for download"
        _set_job(job_id, progress=94, stage=pack_stage)
        zip_path = _make_zip(run_info, work_dir, original_name, model, mode, output_format)
        _set_job(
            job_id,
            status="ready",
            progress=100,
            stage="Stem separation complete",
            zipPath=str(zip_path),
            downloadUrl=f"/jobs/{job_id}/download",
        )
    except subprocess.TimeoutExpired as exc:
        _set_job(
            job_id,
            status="failed",
            progress=100,
            stage="Stem separation timed out",
            error=str(exc),
        )
    except Exception as exc:
        _set_job(
            job_id,
            status="failed",
            progress=100,
            stage="Stem separation failed",
            error=str(exc),
        )


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    return """
<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>Neuralis Stem Worker</title>
  <style>
    body { margin: 0; min-height: 100vh; display: grid; place-items: center; background: #070b0c; color: #eaf8f4; font-family: Arial, sans-serif; }
    main { width: min(760px, calc(100vw - 32px)); border: 1px solid rgba(47, 244, 190, .28); padding: 32px; background: #0b1113; }
    h1 { margin: 0 0 8px; font-size: 24px; letter-spacing: .08em; text-transform: uppercase; }
    p { color: #a7bbb6; line-height: 1.5; }
    label { display: block; margin: 20px 0 8px; font-size: 12px; letter-spacing: .14em; text-transform: uppercase; color: #b9d8d1; }
    input, select { width: 100%; box-sizing: border-box; padding: 12px; color: #fff; border: 1px solid #20343a; background: #070b0c; }
    button { margin-top: 22px; width: 100%; border: 0; padding: 14px; background: #2ff4be; color: #03110e; letter-spacing: .18em; text-transform: uppercase; cursor: pointer; }
    button:disabled { opacity: .55; cursor: wait; }
    .progress { display: none; margin-top: 24px; }
    .track { height: 10px; overflow: hidden; border: 1px solid rgba(47, 244, 190, .25); background: #06100e; }
    .bar { width: 0%; height: 100%; background: linear-gradient(90deg, #2ff4be, #ffd24a); transition: width .35s ease; }
    .status { display: flex; justify-content: space-between; gap: 16px; margin-top: 10px; color: #b9d8d1; font-size: 13px; }
    .download { display: none; margin-top: 18px; color: #2ff4be; letter-spacing: .12em; text-transform: uppercase; }
    .error { display: none; margin-top: 18px; color: #ff8075; line-height: 1.4; }
    code { color: #2ff4be; }
  </style>
</head>
<body>
  <main>
    <h1>Neuralis Stem Worker</h1>
    <p>Private Demucs worker for Neuralis. Use <code>/health</code> for status and <code>/separate</code> for API uploads.</p>
    <form id="stemForm">
      <label>API Key</label>
      <input name="apiKey" type="password" autocomplete="off" />
      <label>Mode</label>
      <select name="mode">
        <option value="fast-2stem" selected>Fast Vocal Enhance - vocals + instrumental</option>
        <option value="premium-4stem">Premium Stem Master - vocals, drums, bass, other</option>
      </select>
      <label>Stem Model</label>
      <select name="model">
        <option value="htdemucs" selected>Demucs Standard - current</option>
        <option value="htdemucs_ft">Demucs Fine-Tuned - higher quality</option>
        <option value="uvr_mdx_voc_ft">UVR MDX Vocal FT - experimental</option>
      </select>
      <label>Output Format</label>
      <select name="format">
        <option value="mp3" selected>MP3 320 kbps - testing</option>
        <option value="wav">WAV - full quality</option>
      </select>
      <label>Audio File</label>
      <input name="file" type="file" accept="audio/*" required />
      <button id="submitButton" type="submit">Separate Stems</button>
    </form>
    <section id="progressPanel" class="progress">
      <div class="track"><div id="progressBar" class="bar"></div></div>
      <div class="status">
        <span id="statusText">Waiting</span>
        <span id="percentText">0%</span>
      </div>
      <a id="downloadLink" class="download" href="#">Download Stems</a>
      <div id="errorText" class="error"></div>
    </section>
  </main>
  <script>
    const form = document.getElementById('stemForm');
    const button = document.getElementById('submitButton');
    const panel = document.getElementById('progressPanel');
    const bar = document.getElementById('progressBar');
    const statusText = document.getElementById('statusText');
    const percentText = document.getElementById('percentText');
    const downloadLink = document.getElementById('downloadLink');
    const errorText = document.getElementById('errorText');

    const setProgress = (value, stage) => {
      const percent = Math.max(0, Math.min(100, Number(value) || 0));
      bar.style.width = `${percent}%`;
      percentText.textContent = `${Math.round(percent)}%`;
      if (stage) statusText.textContent = stage;
    };

    const pollJob = async (id) => {
      const res = await fetch(`/jobs/${id}`);
      const job = await res.json();
      setProgress(job.progress, job.stage || job.status);
      if (job.status === 'ready') {
        button.disabled = false;
        button.textContent = 'Separate Stems';
        downloadLink.href = job.downloadUrl;
        downloadLink.style.display = 'inline-block';
        statusText.textContent = 'Ready';
        return;
      }
      if (job.status === 'failed') {
        button.disabled = false;
        button.textContent = 'Separate Stems';
        errorText.textContent = job.error || 'Stem separation failed';
        errorText.style.display = 'block';
        return;
      }
      setTimeout(() => pollJob(id), 1500);
    };

    form.addEventListener('submit', async (event) => {
      event.preventDefault();
      button.disabled = true;
      button.textContent = 'Processing';
      panel.style.display = 'block';
      downloadLink.style.display = 'none';
      errorText.style.display = 'none';
      setProgress(4, 'Uploading audio');

      const data = new FormData(form);
      const res = await fetch('/jobs', { method: 'POST', body: data });
      const job = await res.json();
      if (!res.ok) {
        button.disabled = false;
        button.textContent = 'Separate Stems';
        errorText.textContent = job.detail || 'Upload failed';
        errorText.style.display = 'block';
        return;
      }
      setProgress(job.progress, job.stage || 'Queued');
      pollJob(job.id);
    });
  </script>
</body>
</html>
"""


@app.get("/health")
def health() -> JSONResponse:
    return JSONResponse(
        {
            "ok": True,
            "service": APP_NAME,
            "model": DEFAULT_MODEL,
            "models": [
                {"id": model, "label": MODEL_LABELS.get(model, model)}
                for model in sorted(ALLOWED_MODELS)
            ],
            "defaultMode": DEFAULT_MODE,
            "defaultFormat": DEFAULT_FORMAT,
            "modes": sorted(ALLOWED_MODES),
            "formats": sorted(ALLOWED_FORMATS),
            "maxUploadMb": MAX_UPLOAD_MB,
            "apiKeyRequired": bool(os.getenv("NEURALIS_STEM_API_KEY", "").strip()),
        }
    )


@app.post("/jobs")
async def create_job(
    request: Request,
    file: UploadFile = File(...),
    model: str = Form(DEFAULT_MODEL),
    mode: str = Form(DEFAULT_MODE),
    format: str = Form(DEFAULT_FORMAT),
    apiKey: str = Form(""),
) -> JSONResponse:
    _check_api_key(request, apiKey)

    selected_model = (model or DEFAULT_MODEL).strip()
    if selected_model not in ALLOWED_MODELS:
        raise HTTPException(status_code=400, detail=f"Unsupported model: {selected_model}")
    selected_mode = (mode or DEFAULT_MODE).strip()
    if selected_mode not in ALLOWED_MODES:
        raise HTTPException(status_code=400, detail=f"Unsupported mode: {selected_mode}")
    selected_format = (format or DEFAULT_FORMAT).strip().lower()
    if selected_format not in ALLOWED_FORMATS:
        raise HTTPException(status_code=400, detail=f"Unsupported format: {selected_format}")

    _cleanup_old_jobs()
    original_name = _safe_name(file.filename)
    suffix = Path(original_name).suffix or ".wav"
    job_id = str(uuid.uuid4())
    work_dir = Path(tempfile.mkdtemp(prefix=f"neuralis-stems-{job_id}-"))

    try:
        input_path = work_dir / f"source{suffix}"
        await _save_upload(file, input_path)
    except Exception:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise

    job = {
        "id": job_id,
        "status": "queued",
        "progress": 10,
        "stage": "Upload received",
        "mode": selected_mode,
        "model": selected_model,
        "format": selected_format,
        "source": original_name,
        "workDir": str(work_dir),
        "createdAt": time.time(),
        "updatedAt": time.time(),
    }
    with JOBS_LOCK:
        JOBS[job_id] = job

    thread = threading.Thread(
        target=_process_job,
        args=(job_id, input_path, work_dir, original_name, selected_model, selected_mode, selected_format),
        daemon=True,
    )
    thread.start()
    return JSONResponse(_public_job(job))


@app.get("/jobs/{job_id}")
def get_job(job_id: str) -> JSONResponse:
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    return JSONResponse(_public_job(job))


@app.get("/jobs/{job_id}/download")
def download_job(job_id: str) -> FileResponse:
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.get("status") != "ready" or not job.get("zipPath"):
        raise HTTPException(status_code=409, detail="Job is not ready")
    zip_path = Path(job["zipPath"])
    if not zip_path.exists():
        raise HTTPException(status_code=404, detail="Stem ZIP was not found")
    return FileResponse(
        zip_path,
        filename=f"neuralis-stems-{job_id}.zip",
        media_type="application/zip",
    )


@app.post("/separate")
async def separate(
    request: Request,
    file: UploadFile = File(...),
    model: str = Form(DEFAULT_MODEL),
    mode: str = Form(DEFAULT_MODE),
    format: str = Form(DEFAULT_FORMAT),
    apiKey: str = Form(""),
) -> FileResponse:
    _check_api_key(request, apiKey)

    selected_model = (model or DEFAULT_MODEL).strip()
    if selected_model not in ALLOWED_MODELS:
        raise HTTPException(status_code=400, detail=f"Unsupported model: {selected_model}")
    selected_mode = (mode or DEFAULT_MODE).strip()
    if selected_mode not in ALLOWED_MODES:
        raise HTTPException(status_code=400, detail=f"Unsupported mode: {selected_mode}")
    selected_format = (format or DEFAULT_FORMAT).strip().lower()
    if selected_format not in ALLOWED_FORMATS:
        raise HTTPException(status_code=400, detail=f"Unsupported format: {selected_format}")

    original_name = _safe_name(file.filename)
    suffix = Path(original_name).suffix or ".wav"
    job_id = str(uuid.uuid4())
    work_dir = Path(tempfile.mkdtemp(prefix=f"neuralis-stems-{job_id}-"))

    try:
        input_path = work_dir / f"source{suffix}"
        await _save_upload(file, input_path)
        run_info = _run_demucs(input_path, work_dir, selected_model, selected_mode)
        zip_path = _make_zip(run_info, work_dir, original_name, selected_model, selected_mode, selected_format)
        return FileResponse(
            zip_path,
            filename=f"neuralis-stems-{job_id}.zip",
            media_type="application/zip",
            background=BackgroundTask(shutil.rmtree, work_dir, ignore_errors=True),
        )
    except HTTPException:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise
    except subprocess.TimeoutExpired as exc:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise HTTPException(status_code=504, detail=f"Stem separation timed out: {exc}") from exc
    except Exception as exc:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise HTTPException(status_code=500, detail=f"Stem separation failed: {exc}") from exc