"""分弦本地后端：音轨分离、试听与下载。"""
import hashlib
import math
import os
import re
import shutil
import subprocess
import threading
import time
import traceback
from urllib.parse import quote

import numpy as np
import soundfile as sf
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

import separate as sep

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))   # app/
# 运行时数据（job、轨、谱面缓存）放在产品目录之外，保持产品目录纯净可核验
DATA = os.environ.get(
    "FENXUAN_DATA",
    os.path.join(os.path.dirname(os.path.dirname(ROOT)),
                 "runtime", "fenxuan", "data"))
os.makedirs(DATA, exist_ok=True)

app = FastAPI(title="分弦")
JOBS = {}          # job_id -> job dict
JOBS_LOCK = threading.Lock()
PROCESS_LOCK = threading.Lock()


@app.middleware("http")
async def no_cache_static(request, call_next):
    """前端文件禁缓存：本地迭代频繁，避免浏览器拿旧版 HTML/JS（曾致按钮失灵）。"""
    resp = await call_next(request)
    path = request.url.path
    if path == "/" or path.endswith((".html", ".js", ".css")):
        resp.headers["Cache-Control"] = "no-cache"
    return resp

STAGE_ORDER = ["decode", "separate", "export"]
MAX_UPLOAD_BYTES = 500 * 1024 * 1024


@app.get("/api/health")
def api_health():
    return {"app": "fenxuan", "status": "ok", "version": "1.5.0"}


def _set(job, stage, frac, msg):
    with JOBS_LOCK:
        job["stage"] = stage
        job["progress"] = round(frac, 3)
        job["message"] = msg


def _fail(job, err):
    traceback.print_exc()
    with JOBS_LOCK:
        job["status"] = "error"
        job["message"] = f"{type(err).__name__}: {err}"


def _run_one(fn, job_id, *args):
    """每次只跑一个重任务，避免多个模型把整台电脑拖慢。"""
    job = JOBS[job_id]
    if not PROCESS_LOCK.acquire(blocking=False):
        _set(job, job.get("stage", "decode"), job.get("progress", 0),
             "前一个任务正在处理，本任务会自动接着开始")
        PROCESS_LOCK.acquire()
    try:
        fn(job_id, *args)
    finally:
        PROCESS_LOCK.release()


def _to_mp3(wav_path, mp3_path):
    if os.path.exists(mp3_path):
        return mp3_path
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-i", wav_path,
         "-codec:a", "libmp3lame", "-b:a", "192k", mp3_path],
        check=True)
    return mp3_path


def _safe_name(value):
    """生成可作为真实文件名的一段文字。"""
    value = re.sub(r'[\\/:*?"<>|\x00-\x1f]', "_", str(value)).strip(" .")
    return value[:80] or "音轨"


def _track_filename(title, *parts, ext):
    return "_".join(_safe_name(x) for x in (title, *parts)) + f".{ext}"


def _signal_stats(path):
    """分块统计信号，避免长音频一次性读入内存。"""
    total_square = 0.0
    total_samples = 0
    windows = []
    with sf.SoundFile(path) as audio:
        for block in audio.blocks(blocksize=max(1, audio.samplerate),
                                  dtype="float32", always_2d=True):
            power = float(np.mean(np.square(block, dtype=np.float64)))
            windows.append(math.sqrt(max(power, 0.0)))
            total_square += power * block.size
            total_samples += block.size
    rms = math.sqrt(total_square / total_samples) if total_samples else 0.0
    return rms, windows


def _quality(path, reference_rms):
    """保守过滤弱轨：只衡量信号可用性，不冒充主观听感评分。"""
    rms, windows = _signal_stats(path)
    relative_db = 20 * math.log10(max(rms, 1e-12) / max(reference_rms, 1e-12))
    threshold = max(10 ** (-45 / 20), reference_rms * 10 ** (-30 / 20))
    active_ratio = (sum(level >= threshold for level in windows) / len(windows)
                    if windows else 0.0)
    # ponytail: 两个环境变量就是听感校准旋钮，有真人盲听数据再换模型。
    try:
        min_relative_db = float(os.environ.get("FENXUAN_MIN_STEM_REL_DB", "-24"))
        min_active_ratio = float(os.environ.get("FENXUAN_MIN_STEM_ACTIVE_RATIO", "0.08"))
    except ValueError:
        min_relative_db, min_active_ratio = -24.0, 0.08
    passed = _passes_quality(relative_db, active_ratio,
                             min_relative_db, min_active_ratio)
    return {"passed": passed, "rmsDbfs": round(20 * math.log10(max(rms, 1e-12)), 2),
            "relativeDb": round(relative_db, 2),
            "activeRatio": round(active_ratio, 3)}


def _passes_quality(relative_db, active_ratio,
                    min_relative_db=-24.0, min_active_ratio=0.08):
    return relative_db >= min_relative_db and active_ratio >= min_active_ratio


OUTPUT_FORMATS = {
    "mp3": ("mp3",),
    "wav": ("wav",),
    "both": ("mp3", "wav"),
}


def _export_audio(wav_path, directory, title, parts, output):
    files = {}
    for ext in OUTPUT_FORMATS[output]:
        filename = _track_filename(title, *parts, ext=ext)
        path = os.path.join(directory, filename)
        if ext == "mp3":
            _to_mp3(wav_path, path)
        elif not os.path.exists(path):
            shutil.copy2(wav_path, path)
        files[ext] = {"name": filename, "path": path}
    return files


def _probe_duration(audio_path):
    """用 ffprobe 读取时长，覆盖页面声明的 MP3/WAV/M4A/FLAC/OGG。"""
    proc = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
        capture_output=True, text=True, check=True)
    duration = float(proc.stdout.strip())
    if duration <= 0:
        raise ValueError("无法读取音频时长")
    return duration


PROFILES = {
    "vocal": {"label": "人声 / 伴奏"},
    "stems6": {"label": "完整乐器分轨"},
}


def run_pipeline(job_id, upload_path, workdir, cache_dir, title,
                 profile="vocal", output="mp3"):
    """纯分轨流水线：解码 → 按档位分离 → 导出音轨。"""
    job = JOBS[job_id]
    try:
        if profile not in PROFILES:
            profile = "vocal"
        cfg = PROFILES[profile]

        _set(job, "decode", 0.03, "解码音频")
        duration = _probe_duration(upload_path)

        rof = {}
        stems = {}
        if profile == "vocal":
            _set(job, "separate", 0.08, "正在分离人声和伴奏")
            rof = sep.separate_roformer(
                upload_path, os.path.join(cache_dir, "roformer"),
                on_progress=lambda f, m: _set(
                    job, "separate", 0.08 + f * 0.82, m))
        else:
            _set(job, "separate", 0.08, "正在一次拆出人声和全部乐器")
            stems = sep.separate(
                upload_path, os.path.join(cache_dir, "stems"),
                on_progress=lambda f, m: _set(
                    job, "separate", 0.08 + f * 0.82, m))

        _set(job, "export", 0.92, "检查清晰度并导出音轨")
        tracks = []
        filtered = []
        track_dir = os.path.join(workdir, "tracks")
        os.makedirs(track_dir, exist_ok=True)
        reference_rms, _ = _signal_stats(upload_path)

        def add_track(name, kind, wav):
            quality = _quality(wav, reference_rms)
            if not quality["passed"]:
                filtered.append({"name": name, "reason": "信号过弱或有效片段过少",
                                 "quality": quality})
                return False
            raw_files = _export_audio(wav, track_dir, title, (name,), output)
            files = {ext: {"name": item["name"],
                           "url": f"/api/tracks/{job_id}/{quote(item['name'])}"}
                     for ext, item in raw_files.items()}
            playback = files.get("mp3") or files["wav"]
            tracks.append({"name": name, "kind": kind, "url": playback["url"],
                           "files": files, "quality": quality})
            return True

        names = {"drums": "鼓", "bass": "贝斯", "guitar": "吉他",
                 "piano": "钢琴", "other": "其他器乐"}

        if profile == "vocal":
            add_track("人声", "melody", rof["vocals"])
            add_track("伴奏", "accompaniment", rof["instrumental"])
        else:
            add_track("人声", "melody", stems["vocals"])
            for stem, name in names.items():
                if stem in stems:
                    add_track(name, "stem", stems[stem])

        keep = {item["name"] for track in tracks for item in track["files"].values()}
        for filename in os.listdir(track_dir):
            if filename.lower().endswith((".mp3", ".wav")) and filename not in keep:
                os.remove(os.path.join(track_dir, filename))

        with JOBS_LOCK:
            job.update({
                "status": "done", "stage": "export", "progress": 1.0,
                "message": "完成" if tracks else "完成，未找到达到清晰标准的音轨",
                "duration": duration, "title": title, "profile": profile,
                "output": output,
            })
            melody = next((t["url"] for t in tracks if t["kind"] == "melody"), "")
            accompaniment = next(
                (t["url"] for t in tracks if t["kind"] == "accompaniment"), "")
            job["result"] = {
                "duration": round(duration, 2),
                "profile": profile,
                "profileLabel": cfg["label"],
                "output": output,
                "tracks": tracks,
                "filteredTracks": filtered,
                "audio": {
                    "melody": melody,
                    "accompaniment": accompaniment,
                },
            }
    except Exception as e:
        _fail(job, e)


# ---------------- API ----------------

@app.post("/api/analyze")
async def api_analyze(file: UploadFile = File(...), profile: str = Form("vocal"),
                      output: str = Form("mp3")):
    raw = await file.read(MAX_UPLOAD_BYTES + 1)
    if len(raw) < 1000:
        raise HTTPException(400, "文件太小，不是有效音频")
    if len(raw) > MAX_UPLOAD_BYTES:
        raise HTTPException(413, "音频超过 500MB，请先压缩或截取需要处理的部分")
    if profile not in PROFILES:
        raise HTTPException(400, "未知的分轨方案")
    if output not in OUTPUT_FORMATS:
        raise HTTPException(400, "未知的输出格式")
    safe_name = os.path.basename(file.filename or "audio.wav")
    cache_id = hashlib.sha1(raw + profile.encode()).hexdigest()[:16]
    job_id = hashlib.sha1(
        raw + profile.encode() + output.encode() + safe_name.encode()).hexdigest()[:16]
    with JOBS_LOCK:
        existing = JOBS.get(job_id)
        if existing and existing["status"] == "done":
            return {"job": job_id, "cached": True}
        if existing and existing["status"] == "running":
            return {"job": job_id, "running": True}  # 重连即可，不重复跑
    workdir = os.path.join(DATA, job_id)
    cache_dir = os.path.join(DATA, cache_id)
    os.makedirs(workdir, exist_ok=True)
    os.makedirs(cache_dir, exist_ok=True)
    ext = os.path.splitext(safe_name)[1] or ".wav"
    upload_path = os.path.join(workdir, f"original{ext}")
    with open(upload_path, "wb") as f:
        f.write(raw)
    try:
        _probe_duration(upload_path)
    except (OSError, ValueError, subprocess.SubprocessError):
        shutil.rmtree(workdir, ignore_errors=True)
        raise HTTPException(400, "无法读取这个音频，请换用 MP3、WAV、M4A、FLAC 或 OGG")
    title = os.path.splitext(safe_name)[0]
    with JOBS_LOCK:
        JOBS[job_id] = {"status": "running", "stage": "decode",
                        "progress": 0.0, "message": "排队中",
                        "workdir": workdir, "created": time.time()}
    threading.Thread(target=_run_one,
                     args=(run_pipeline, job_id, upload_path, workdir, cache_dir,
                           title, profile, output),
                     daemon=True).start()
    return {"job": job_id, "cached": False}


@app.get("/api/jobs/{job_id}")
def api_job(job_id: str):
    job = JOBS.get(job_id)
    if not job:
        raise HTTPException(404, "job 不存在（服务可能已重启）")
    out = {k: job.get(k) for k in
           ("status", "stage", "progress", "message", "result")}
    out["stageIndex"] = STAGE_ORDER.index(job["stage"]) if job.get("stage") in STAGE_ORDER else 0
    return out


@app.get("/api/tracks/{job_id}/{name}")
def api_tracks(job_id: str, name: str):
    job = JOBS.get(job_id)
    ext = os.path.splitext(name)[1].lower()
    if not job or ext not in (".mp3", ".wav") or "/" in name or ".." in name:
        raise HTTPException(404, "音轨不存在")
    path = os.path.join(job["workdir"], "tracks", name)
    if not os.path.exists(path):
        raise HTTPException(404, "音轨不存在")
    return FileResponse(path, media_type="audio/mpeg" if ext == ".mp3" else "audio/wav",
                        filename=name)


@app.get("/api/audio/{job_id}/{name}")
def api_audio(job_id: str, name: str):
    # 兼容旧路径：重定向到分轨
    job = JOBS.get(job_id)
    if not job:
        raise HTTPException(404, "音频不存在")
    kind = "melody" if name == "melody.mp3" else "accompaniment"
    track = next((t for t in job.get("result", {}).get("tracks", [])
                  if t["kind"] == kind), None)
    if not track:
        raise HTTPException(404, "音频不存在")
    item = track["files"].get("mp3") or track["files"].get("wav")
    path = os.path.join(job["workdir"], "tracks", item["name"])
    if not os.path.exists(path):
        raise HTTPException(404, "音频不存在")
    return FileResponse(path)


class ExportReq(BaseModel):
    job: str
    tracks: list[int]          # 要打包的轨道序号，至少一条


@app.post("/api/export-zip")
def api_export_zip(req: ExportReq):
    """按用户明确勾选的轨道打包 ZIP。"""
    job = JOBS.get(req.job)
    if not job or job.get("status") != "done":
        raise HTTPException(409, "分析尚未完成或 job 不存在")
    import zipfile, io
    result = job["result"]
    tracks = result["tracks"]
    if not req.tracks:
        raise HTTPException(400, "请至少选择一条音轨")
    idxs = list(dict.fromkeys(req.tracks))
    if any(i < 0 or i >= len(tracks) for i in idxs):
        raise HTTPException(400, "选择的音轨不存在")
    track_dir = os.path.join(job["workdir"], "tracks")
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
        for i in idxs:
            for item in tracks[i]["files"].values():
                z.write(os.path.join(track_dir, item["name"]), item["name"])
    buf.seek(0)
    from fastapi.responses import StreamingResponse
    from urllib.parse import quote
    fname = quote(f"{job.get('title','tracks')}_分轨.zip")
    return StreamingResponse(
        buf, media_type="application/zip",
        headers={"Content-Disposition":
                 f"attachment; filename=tracks.zip; filename*=UTF-8''{fname}"})


# ---------------- 静态前端 ----------------

app.mount("/", StaticFiles(directory=ROOT, html=True), name="static")


def _reaper():
    while True:
        time.sleep(600)
        now = time.time()
        with JOBS_LOCK:
            for jid, job in list(JOBS.items()):
                if now - job.get("created", now) > 2 * 3600:
                    shutil.rmtree(job.get("workdir", ""), ignore_errors=True)
                    del JOBS[jid]


if __name__ == "__main__":
    import uvicorn
    threading.Thread(target=_reaper, daemon=True).start()
    uvicorn.run(app, host="127.0.0.1", port=8642, log_level="warning")
