# robust_relay.py
# 高稳定性 YouTube 转播模块（防误判下播 + 看门狗 + 自动重连 + token刷新）
# 作者：融合 A+B 代码精华，专为抖音/抖音国际版长期无人值守转播设计

import subprocess
import time
import threading
import os
import signal
import datetime
from typing import List, Optional, Callable

# ==================== 配置常量（可全局覆盖） ====================
WATCHDOG_SILENCE = 600        # 10分钟无输出 → 软重启
WATCHDOG_GRACE = 120          # 启动宽限期
REFRESH_URL_EVERY = 300       # 每5分钟强制刷新 real_url（防token过期）
OFFLINE_DEBOUNCE = 3          # 连续3次API返回下播才认为真下播
OFFLINE_WINDOW = 120          # 最近120秒有活跃就忽略下播信号
SOFT_RESTART_COOLDOWN = 90    # 软重启最小间隔

# ==================== 轻量探活 ====================
def _probe_alive(url: str, timeout: float = 5.0) -> bool:
    if not url:
        return False
    try:
        import urllib.request, contextlib
        opener = urllib.request.build_opener()
        if hasattr(urllib.request, 'getproxies'):
            proxy = urllib.request.getproxies().get('http')
            if proxy:
                opener.add_handler(urllib.request.ProxyHandler({"http": proxy, "https": proxy}))
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"}, method="HEAD")
        with contextlib.closing(opener.open(req, timeout=timeout)) as r:
            if r.code >= 400:
                return False
            head = r.read(256)
            if url.lower().endswith(".m3u8"):
                return b"#EXTM3U" in head.upper()
            if url.lower().endswith(".flv"):
                return head.startswith(b"FLV")
            return len(head) > 0
    except:
        return False

# ==================== 增强版 ffmpeg 启动器 ====================
def _launch_ffmpeg(
    input_url: str,
    youtube_rtmp: str,
    anchor_name: str,
    logger: Callable[[str], None],
    extra_cmd: List[str] = None
) -> subprocess.Popen:

    cmd = [
        "ffmpeg", "-nostdin", "-hide_banner", "-stats", "-stats_period", "60", "-loglevel", "info",

        # 输入端超强重连
        "-reconnect", "1", "-reconnect_streamed", "1", "-reconnect_at_eof", "1",
        "-reconnect_delay_max", "30", "-rw_timeout", "20000000",

        # 极致低延迟 + 时间戳修复
        "-fflags", "+genpts+igndts+flush_packets", "-flags", "+global_header",
        "-thread_queue_size", "4096", "-analyzeduration", "0", "-probesize", "32",

        "-i", input_url,

        # copy 模式（质量无损）
        "-c:v", "copy", "-c:a", "copy", "-bsf:a", "aac_adtstoasc",
        "-map", "0:v?", "-map", "0:a?",

        # YouTube 强制参数
        "-f", "mpegts", "-g", "2", "-rtmp_buffer", "1000", "-rtmp_live", "live",
        "-max_delay", "0", "-tune", "zerolatency", "-avoid_negative_ts", "make_zero",
        "-muxdelay", "0", "-muxpreload", "0",
        "-fflags", "+nobuffer+flush_packets", "-avioflags", "direct",
    ]

    if extra_cmd:
        cmd.extend(extra_cmd)

    cmd.append(youtube_rtmp)

    creationflags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=0,
        creationflags=creationflags
    )

    # 用于看门狗
    proc.last_output_ts = time.time()

    def _reader():
        try:
            for line in proc.stdout:  # type: ignore
                if not line:
                    break
                text = line.decode(errors="ignore").rstrip()
                proc.last_output_ts = time.time()
                if any(k in text for k in ["bitrate=", "frame=", "speed=", "fps="]):
                    logger(f"[{anchor_name}] {text}")
        except:
            pass
    threading.Thread(target=_reader, daemon=True).start()

    return proc

# ==================== 主函数：高稳定性转播（替代原 check_subprocess） ====================
def robust_youtube_relay(
    anchor_name: str,
    real_url: str,
    youtube_rtmp: str,
    logger: Callable[[str], None],
    get_latest_stream_url: Callable[[], Optional[str]],  # 用于刷新 real_url（如有token过期）
    is_stream_still_live: Callable[[], bool],            # 检查主播是否在播（API）
    extra_ffmpeg_args: List[str] = None,                 # 可选：加 -f segment 分段保存
    heartbeat_callback: Callable[[int], None] = None     # 可选：每分钟回调已运行秒数
) -> bool:
    """
    返回值：
        True  = 主播已下播（正常结束）
        False = 发生未知错误（建议计数错误次数）
    """
    proc: Optional[subprocess.Popen] = None
    start_time = time.time()
    last_refresh = 0.0
    last_soft_restart = 0.0
    offline_count = 0
    last_live_ts = time.time()

    try:
        proc = _launch_ffmpeg(real_url, youtube_rtmp, anchor_name, logger, extra_ffmpeg_args)
        logger(f"[{anchor_name}] 已开始高稳定性转播 → {youtube_rtmp.split('/')[-1]}")

        while True:
            now = time.time()
            elapsed = int(now - start_time)
            silence = now - proc.last_output_ts if proc else 0

            # 1. 进程意外退出 → 重启
            if proc and proc.poll() is not None:
                logger(f"[{anchor_name}] ffmpeg 意外退出，3秒后重启...")
                time.sleep(3)
                proc = _launch_ffmpeg(real_url, youtube_rtmp, anchor_name, logger, extra_ffmpeg_args)
                start_time = time.time()
                continue

            # 2. 看门狗：长时间无输出 → 软重启
            if elapsed > WATCHDOG_GRACE and silence > WATCHDOG_SILENCE:
                if now - last_soft_restart > SOFT_RESTART_COOLDOWN:
                    logger(f"[{anchor_name}] 检测到卡顿（{silence:.0f}s无输出），软重启 ffmpeg...")
                    last_soft_restart = now
                    if proc and proc.stdin:
                        proc.stdin.write(b'q')
                        proc.stdin.flush()
                    time.sleep(4)
                    proc = _launch_ffmpeg(real_url, youtube_rtmp, anchor_name, logger, extra_ffmpeg_args)
                    start_time = time.time()
                    continue

            # 3. 定期刷新播放地址（防token过期）
            if now - last_refresh > REFRESH_URL_EVERY:
                new_url = get_latest_stream_url()
                if new_url and new_url != real_url:
                    logger(f"[{anchor_name}] 检测到新的播放地址，切换中...")
                    real_url = new_url
                    if proc and proc.stdin:
                        proc.stdin.write(b'q')
                    time.sleep(3)
                    proc = _launch_ffmpeg(real_url, youtube_rtmp, anchor_name, logger, extra_ffmpeg_args)
                    start_time = time.time()
                last_refresh = now

            # 4. 下播防抖判定（核心防误判）
            try:
                still_live = is_stream_still_live()
            except:
                still_live = True  # 接口异常不判下播

            if not still_live:
                # 交叉验证：流还能访问 → 忽略API误判
                if _probe_alive(real_url):
                    logger(f"[{anchor_name}] API判下播但流仍活跃，忽略本次误判")
                    still_live = True
                    last_live_ts = now

            if still_live:
                offline_count = 0
                last_live_ts = now
            else:
                if now - last_live_ts < OFFLINE_WINDOW:
                    logger(f"[{anchor_name}] 短期掉线，忽略")
                else:
                    offline_count += 1
                    if offline_count >= OFFLINE_DEBOUNCE:
                        logger(f"[{anchor_name}] 连续{offline_count}次确认下播，结束转播")
                        if proc and proc.stdin:
                            proc.stdin.write(b'q')
                        return True
                    else:
                        logger(f"[{anchor_name}] 疑似下播 {offline_count}/{OFFLINE_DEBOUNCE}，继续观察")

            # 5. 心跳日志 & 回调
            if elapsed % 60 < 2:
                h, m = divmod(elapsed // 60, 60)
                logger(f"[{anchor_name}] 正在转播 {h:02d}:{m:02d}")
                if heartbeat_callback:
                    heartbeat_callback(elapsed)

            time.sleep(8)

    except KeyboardInterrupt:
        logger(f"[{anchor_name}] 手动终止转播")
        return True
    except Exception as e:
        logger(f"[{anchor_name}] 严重错误: {e}")
        return False
    finally:
        if proc:
            try:
                proc.stdin.write(b'q') if proc.stdin else None
                proc.wait(timeout=8)
            except:
                proc.kill()