"""分弦本地后端：音轨分离、试听与下载。"""
import hashlib
import math
import os
import re
import shutil
import subprocess
import tempfile
import threading
import time
import traceback
from contextlib import ExitStack
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.6.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 _ffmpeg():
    system = shutil.which("ffmpeg")
    if system:
        return system
    from imageio_ffmpeg import get_ffmpeg_exe
    return get_ffmpeg_exe()


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):
    """优先直接读时长；M4A 等格式再由内置 ffmpeg 快速探测。"""
    try:
        info = sf.info(audio_path)
        duration = info.frames / info.samplerate
    except (RuntimeError, ValueError):
        proc = subprocess.run(
            [_ffmpeg(), "-hide_banner", "-i", audio_path],
            capture_output=True, text=True)
        match = re.search(r"Duration: (\d+):(\d+):(\d+(?:\.\d+)?)", proc.stderr)
        if not match:
            raise ValueError("无法读取音频时长")
        hours, minutes, seconds = match.groups()
        duration = int(hours) * 3600 + int(minutes) * 60 + float(seconds)
    if duration <= 0:
        raise ValueError("无法读取音频时长")
    return duration


def _mix_stems(paths, output_path):
    """分块合成伴奏；不再运行第二个模型，也不一次读完整首歌。"""
    if os.path.exists(output_path):
        return output_path
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with ExitStack() as stack:
        sources = [stack.enter_context(sf.SoundFile(path)) for path in paths]
        first = sources[0]
        if any((src.samplerate, src.channels) !=
               (first.samplerate, first.channels) for src in sources[1:]):
            raise ValueError("音轨采样格式不一致，无法合成伴奏")
        frames = min(len(src) for src in sources)
        with sf.SoundFile(output_path, "w", first.samplerate, first.channels,
                          subtype="FLOAT") as target:
            written = 0
            while written < frames:
                count = min(first.samplerate, frames - written)
                mix = sum((src.read(count, dtype="float32", always_2d=True)
                           for src in sources),
                          np.zeros((count, first.channels), dtype="float32"))
                target.write(mix)
                written += count
    return output_path


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)

        _set(job, "separate", 0.08, "正在一次拆出人声和全部乐器")
        stem_cache = os.path.join(cache_dir, "stems")
        stems = sep.separate(
            upload_path, stem_cache,
            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":
            instrumental = _mix_stems(
                [stems[name] for name in names],
                os.path.join(stem_cache, "instrumental.wav"))
            add_track("人声", "melody", stems["vocals"])
            add_track("伴奏", "accompaniment", 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")):
    if profile not in PROFILES:
        raise HTTPException(400, "未知的分轨方案")
    if output not in OUTPUT_FORMATS:
        raise HTTPException(400, "未知的输出格式")

    incoming = os.path.join(DATA, "incoming")
    os.makedirs(incoming, exist_ok=True)
    digest = hashlib.sha1()
    total = 0
    temp_path = ""
    try:
        with tempfile.NamedTemporaryFile(dir=incoming, delete=False) as target:
            temp_path = target.name
            while chunk := await file.read(1024 * 1024):
                total += len(chunk)
                if total > MAX_UPLOAD_BYTES:
                    raise HTTPException(
                        413, "音频超过 500MB，请先压缩或截取需要处理的部分")
                digest.update(chunk)
                target.write(chunk)
        if total < 1000:
            raise HTTPException(400, "文件太小，不是有效音频")
    except Exception:
        if temp_path:
            os.remove(temp_path)
        raise

    safe_name = os.path.basename(file.filename or "audio.wav")
    content_id = digest.hexdigest()[:16]
    cache_id = content_id
    job_id = hashlib.sha1(
        (content_id + profile + output + safe_name).encode()).hexdigest()[:16]
    with JOBS_LOCK:
        existing = JOBS.get(job_id)
        if existing and existing["status"] == "done":
            os.remove(temp_path)
            return {"job": job_id, "cached": True}
        if existing and existing["status"] == "running":
            os.remove(temp_path)
            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}")
    os.replace(temp_path, upload_path)
    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, "cache_dir": cache_dir,
                        "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 _cleanup_storage(now=None):
    """清除服务重启后遗留的旧任务；运行中的目录永远跳过。"""
    now = now or time.time()
    with JOBS_LOCK:
        active = {job.get(key) for job in JOBS.values()
                  if job.get("status") == "running"
                  for key in ("workdir", "cache_dir")}
    for entry in os.scandir(DATA):
        if entry.path in active:
            continue
        if entry.name == "incoming" and entry.is_dir():
            for pending in os.scandir(entry.path):
                if now - pending.stat().st_mtime > 2 * 3600:
                    os.remove(pending.path)
            continue
        if entry.is_dir() and now - entry.stat().st_mtime > 2 * 3600:
            shutil.rmtree(entry.path, ignore_errors=True)


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]
        _cleanup_storage(now)


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