# demucs_pipeline.py
# Demucs 人声分离管道：FFmpeg 拉流 → Demucs 分离人声 → FFmpeg 推流
# 兼容 Windows / Ubuntu，支持 NVIDIA GPU 加速与 CPU 回退
"""
Demucs 音频管道模块
用于将混音分离为人声，推送到 YouTube 时避免 Content ID 识别背景音乐。
"""
import os
import sys
import socket
import threading
import subprocess
import queue
import time
import platform

# 可选依赖：Demucs 和 PyTorch
DEMUCS_AVAILABLE = False
DEMUCS_MODEL = None
TORCH_AVAILABLE = False

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    pass

try:
    from demucs.pretrained import get_model
    from demucs.apply import apply_model
    DEMUCS_AVAILABLE = True
except ImportError:
    pass


def _check_demucs_available():
    """检查 Demucs 是否可用"""
    return DEMUCS_AVAILABLE and TORCH_AVAILABLE


def _load_demucs_model(use_gpu: bool):
    """加载 Demucs 模型，支持 GPU/CPU 回退"""
    global DEMUCS_MODEL
    if not _check_demucs_available():
        return None
    try:
        # htdemucs 较轻量，适合实时；htdemucs_ft 质量更好但更慢
        model_name = "htdemucs"
        device = "cpu"
        if use_gpu and TORCH_AVAILABLE and torch.cuda.is_available():
            try:
                device = "cuda"
            except Exception:
                device = "cpu"
        model = get_model(model_name)
        model.eval()
        return model.to(device)
    except Exception as e:
        print(f"[WARN] Demucs 模型加载失败，回退原始音频: {e}")
        return None


# Demucs 模型采样率
DEMUCS_SAMPLE_RATE = 44100
DEMUCS_CHANNELS = 2
# 每次处理的音频长度（秒），影响延迟与质量
DEMUCS_CHUNK_SEC = 2.0
# 字节 per sample (s16le: 2 bytes)
BYTES_PER_SAMPLE = 2


def _demucs_separate_vocals(model, audio_chunk, device):
    """对音频块进行人声分离，返回人声轨道"""
    import torch
    # audio_chunk: numpy or tensor, shape (channels, samples)
    if hasattr(audio_chunk, 'numpy'):
        wav = audio_chunk
    else:
        import numpy as np
        wav = torch.from_numpy(audio_chunk).float()
    # Demucs 期望 (batch, channels, samples)
    if wav.dim() == 2:
        wav = wav.unsqueeze(0)
    with torch.no_grad():
        out = apply_model(model, wav, device=device, progress=False, overlap=0.25)
    # out: (batch, sources, channels, samples)
    # sources: [drums, bass, other, vocals]
    vocals = out[0, 3]  # 取 vocals
    return vocals.cpu().numpy()


def _run_demucs_worker(audio_read_pipe, audio_write_socket, model, device, stop_event):
    """Demucs 工作线程：从 FFmpeg 读音频，分离人声，写入 socket"""
    import numpy as np
    chunk_samples = int(DEMUCS_SAMPLE_RATE * DEMUCS_CHUNK_SEC * DEMUCS_CHANNELS)
    chunk_bytes = chunk_samples * BYTES_PER_SAMPLE
    buf = b""

    def read_audio(size):
        if hasattr(audio_read_pipe, 'recv'):
            return audio_read_pipe.recv(size)
        return audio_read_pipe.read(size)

    try:
        while not stop_event.is_set():
            try:
                data = read_audio(65536)
                if not data:
                    break
                buf += data
            except (socket.error, OSError, ValueError):
                break
            while len(buf) >= chunk_bytes:
                chunk = buf[:chunk_bytes]
                buf = buf[chunk_bytes:]
                # s16le -> float
                arr = np.frombuffer(chunk, dtype=np.int16)
                arr = arr.reshape(-1, DEMUCS_CHANNELS).T.astype(np.float32) / 32768.0
                try:
                    vocals = _demucs_separate_vocals(model, arr, device)
                    # float -> s16le
                    out = (vocals * 32768).astype(np.int16).tobytes()
                    try:
                        audio_write_socket.sendall(out)
                    except (socket.error, OSError):
                        return
                except Exception as e:
                    print(f"[WARN] Demucs 处理异常: {e}")
                    # 回退：直接透传原始音频
                    try:
                        audio_write_socket.sendall(chunk)
                    except (socket.error, OSError):
                        return
        # 剩余 buffer
        if buf and audio_write_socket:
            try:
                arr = np.frombuffer(buf, dtype=np.int16)
                arr = arr.reshape(-1, DEMUCS_CHANNELS).T.astype(np.float32) / 32768.0
                vocals = _demucs_separate_vocals(model, arr, device)
                out = (vocals * 32768).astype(np.int16).tobytes()
                audio_write_socket.sendall(out)
            except Exception:
                pass
    except Exception as e:
        print(f"[WARN] Demucs worker 异常: {e}")
    finally:
        try:
            audio_write_socket.shutdown(socket.SHUT_WR)
            audio_write_socket.close()
        except Exception:
            pass


def run_demucs_pipeline(
    real_url: str,
    youtube_rtmp: str,
    douyin_headers: str,
    nvenc_params: dict,
    end_keywords: list,
    stderr_queue_timeout: float,
    no_frame_timeout: int,
    start_check_after: int,
    stale_output_seconds: int,
    stream_ended_exception,
    cleanup_proc_fn,
    stderr_reader_fn,
):
    """
    运行 Demucs 管道：拉流 → 人声分离 → 推流
    若主播下播应停止则 raise stream_ended_exception
    """

    # 检查 Demucs
    use_gpu = nvenc_params.get("demucs_use_gpu", True)
    model = _load_demucs_model(use_gpu)
    if model is None:
        raise RuntimeError("Demucs 模型不可用，将回退原始推流")

    device = next(model.parameters()).device
    device_str = "cuda" if "cuda" in str(device) else "cpu"
    print(f"[INFO] Demucs 人声分离已启用 ({device_str})")

    # TCP 服务器用于音频（跨平台）
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(("127.0.0.1", 0))
    port = server.getsockname()[1]
    server.listen(1)

    tcp_audio_url = f"tcp://127.0.0.1:{port}"
    stop_event = threading.Event()
    conn = [None]  # 用 list 以便在闭包中修改

    def accept_conn():
        server.settimeout(30)
        try:
            c, _ = server.accept()
            conn[0] = c
        except Exception:
            pass
        finally:
            try:
                server.close()
            except Exception:
                pass

    accept_thread = threading.Thread(target=accept_conn, daemon=True)
    accept_thread.start()
    time.sleep(0.2)  # 确保 server 已 listen

    # FFmpeg 拉流（视频）
    ffmpeg_v_cmd = [
        "ffmpeg", "-y", "-loglevel", "info", "-nostdin", "-re",
        "-stats", "-stats_period", "1",
        "-rw_timeout", "30000000",
        "-reconnect", "1", "-reconnect_at_eof", "1", "-reconnect_streamed", "1",
        "-reconnect_delay_max", "5",
        "-fflags", "+genpts+discardcorrupt+nobuffer+flush_packets",
        "-err_detect", "ignore_err",
        "-headers", douyin_headers,
        "-i", real_url,
        "-map", "0:v", "-c:v", "copy", "-f", "nut", "pipe:1",
    ]

    # FFmpeg 拉流（音频）
    ffmpeg_a_cmd = [
        "ffmpeg", "-y", "-loglevel", "error", "-nostdin", "-re",
        "-rw_timeout", "30000000",
        "-reconnect", "1", "-reconnect_at_eof", "1", "-reconnect_streamed", "1",
        "-headers", douyin_headers,
        "-i", real_url,
        "-map", "0:a", "-f", "s16le", "-ac", str(DEMUCS_CHANNELS),
        "-ar", str(DEMUCS_SAMPLE_RATE), "pipe:1",
    ]

    # FFmpeg 推流：音频 tcp 放第一（FFmpeg 按顺序打开输入，pipe:0 会阻塞直到有数据，
    # 若 pipe 在前则永远连不上 tcp，导致 accept 超时）
    vf = "fps=30,scale=720:1280:force_original_aspect_ratio=decrease:flags=lanczos,pad=720:1280:(ow-iw)/2:(oh-ih)/2:black"
    ffmpeg_out_cmd = [
        "ffmpeg", "-y", "-loglevel", "info", "-nostdin",
        "-f", "s16le", "-ac", str(DEMUCS_CHANNELS), "-ar", str(DEMUCS_SAMPLE_RATE),
        "-i", tcp_audio_url,
        "-f", "nut", "-i", "pipe:0",
        "-map", "1:v", "-map", "0:a",
        "-vf", vf,
        "-c:v", nvenc_params.get("codec_v", "libx264"),
        "-preset", nvenc_params.get("preset_v", "fast"),
        "-profile:v", "high",
        "-b:v", "2600k", "-maxrate", "3000k", "-bufsize", "5200k",
        "-vsync", "cfr", "-g", "60", "-keyint_min", "60", "-r", "30",
        "-bf", "0", "-sc_threshold", "0",
        "-x264-params", "force-cfr=1:scenecut=0:open_gop=0:sync-lookahead=0",
        "-pix_fmt", "yuv420p",
        "-c:a", "aac", "-b:a", "128k", "-ar", "48000", "-ac", "2",
        "-af", "aresample=async=1:first_pts=0",
        "-flags", "+global_header", "-flvflags", "no_duration_filesize",
        "-f", "flv", youtube_rtmp,
    ]
    if nvenc_params.get("tune_param"):
        idx = ffmpeg_out_cmd.index("-profile:v")
        for p in reversed(nvenc_params["tune_param"]):
            ffmpeg_out_cmd.insert(idx, p)

    proc_v = None
    proc_a = None
    proc_out = None
    demucs_thread = None

    try:
        proc_v = subprocess.Popen(
            ffmpeg_v_cmd,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=1024 * 1024,
        )
        proc_a = subprocess.Popen(
            ffmpeg_a_cmd,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            bufsize=1024 * 1024,
        )
        proc_out = subprocess.Popen(
            ffmpeg_out_cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            bufsize=1024 * 1024,
        )

        # 等待 TCP 连接（FFmpeg_out 会连接）
        accept_thread.join(timeout=15)
        if conn[0] is None:
            raise RuntimeError("Demucs 管道: 音频 TCP 连接超时")

        # 视频转发线程
        def copy_video():
            try:
                while proc_v.poll() is None and proc_out.poll() is None:
                    chunk = proc_v.stdout.read(1024 * 256)
                    if not chunk:
                        break
                    proc_out.stdin.write(chunk)
            except Exception:
                pass
            finally:
                try:
                    proc_out.stdin.close()
                except Exception:
                    pass

        copy_t = threading.Thread(target=copy_video, daemon=True)
        copy_t.start()

        # Demucs 工作线程（从 proc_a 读，处理后写 conn）
        # proc_a 输出是 stdout，conn 是 socket。需要把 proc_a.stdout 的数据通过 socket 发到 FFmpeg。
        # FFmpeg_out 的第二个输入是 tcp，已经连接了 conn。所以我们从 proc_a 读，Demucs 处理，写 conn。
        # 但 proc_a.stdout 是 pipe，conn 是 socket。我们需要一个线程读 proc_a 并写 conn。
        # 注意：proc_a 输出到 stdout，我们读。conn 是 FFmpeg 作为 client 连上来的，所以 FFmpeg 在读，我们写 conn。
        demucs_thread = threading.Thread(
            target=_run_demucs_worker,
            args=(proc_a.stdout, conn[0], model, device, stop_event),
            daemon=True,
        )
        demucs_thread.start()

        # 监控 proc_v 的 stderr（与源连接，用于检测下播）
        stderr_q = queue.Queue()
        reader_t = threading.Thread(target=stderr_reader_fn, args=(proc_v, stderr_q))
        reader_t.daemon = True
        reader_t.start()

        # 监控 proc_out 的 stderr（推流到 YouTube，错误必须可见）
        def _proc_out_stderr_reader():
            try:
                for line in iter(proc_out.stderr.readline, b""):
                    if line:
                        decoded = line.decode("utf-8", errors="replace").strip()
                        if decoded:
                            print(f"[YouTube推流] {decoded}")
            except Exception:
                pass

        proc_out_reader_t = threading.Thread(target=_proc_out_stderr_reader, daemon=True)
        proc_out_reader_t.start()

        start_time = last_frame_time = time.time()
        proc_v.last_out_ts = time.time()

        while True:
            try:
                line = stderr_q.get(timeout=stderr_queue_timeout)
            except queue.Empty:
                # 若 proc_out（推流进程）已异常退出，立即报告
                if proc_out.poll() is not None and proc_out.returncode != 0:
                    print(f"[ERROR] YouTube 推流进程已退出，返回码: {proc_out.returncode}")
                    break
                if time.time() - getattr(proc_v, "last_out_ts", time.time()) > stale_output_seconds:
                    print("[WARN] 长时间无任何日志，强制重启进程")
                    break
                continue

            if line is None:
                if proc_v.poll() is not None:
                    break
                continue

            # proc_v 的 stderr 为 bytes，需解码为 str
            if isinstance(line, bytes):
                line = line.decode("utf-8", errors="replace")
            line = line.strip()
            if line.startswith("__ERR__READER__"):
                print(f"[ERROR] stderr reader 异常: {line}")
                continue

            if line:
                proc_v.last_out_ts = time.time()
                print(f"[FFmpeg] {line}")
                line_lower = line.lower()
                if any(kw in line_lower for kw in end_keywords):
                    print("[ERROR] 检测到主播真下播或被明确拒绝，停止推流")
                    raise stream_ended_exception()
                if "frame=" in line_lower or "fps=" in line_lower:
                    last_frame_time = time.time()
                if time.time() - start_time > start_check_after:
                    if time.time() - last_frame_time > no_frame_timeout:
                        print(f"[WARN] 连续 {no_frame_timeout}s 无新帧，判定主播已下播")
                        raise stream_ended_exception()

    except stream_ended_exception:
        raise
    except Exception as e:
        print(f"[ERROR] Demucs 管道异常: {e}")
        import traceback
        traceback.print_exc()
    finally:
        stop_event.set()
        if proc_out is not None and proc_out.poll() is not None and proc_out.returncode != 0:
            print(f"[INFO] YouTube 推流进程退出码: {proc_out.returncode}")
        cleanup_proc_fn(proc_v)
        cleanup_proc_fn(proc_a)
        cleanup_proc_fn(proc_out)
        try:
            if conn[0]:
                conn[0].close()
        except Exception:
            pass
        try:
            server.close()
        except Exception:
            pass