#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Douyin → YouTube 极致稳定推流器 (async 版)
- 特性：
  * 全异步 asyncio 架构，单线程高并发 I/O；
  * 彩色/文件双日志，结构化前缀、滚动保存；
  * FFMPEG 自恢复与指数退避 + 抖动；
  * HLS/HTTP 强化：-reconnect-* 自动断线重连，低延迟参数；
  * 可选本地录制（独立进程，意外退出自动重启）；
  * 优雅退出：SIGINT/SIGTERM，确保子进程回收；
  * list.ini 热加载（mtime 变更即时生效）；
  * 健康状态心跳（码率/时长/是否已推上帧）；
  * 配置集中化：config.ini + 环境变量覆盖；

依赖：仅标准库 + colorama + aiohttp
保留你现有 spider/stream 模块调用约定。
"""

from __future__ import annotations
import asyncio
import contextlib
import dataclasses
import os
import signal
import sys
import time
import uuid
import logging
import configparser
from typing import Optional, Tuple, List

import aiohttp
from colorama import init as color_init, Fore, Style

# ====== 第三方模块日志降噪（尽量提前） ======
for noisy in ("aiohttp", "aiohttp.client", "httpcore", "urllib3", "asyncio"):
    logging.getLogger(noisy).setLevel(logging.CRITICAL)

# 你本地模块（按需降噪）
import spider  # 保持与现有工程兼容
import stream  # 保持与现有工程兼容
if hasattr(spider, 'logger'):
    spider.logger.setLevel(logging.CRITICAL)
if hasattr(stream, 'logger'):
    stream.logger.setLevel(logging.CRITICAL)

# ================= 日志系统 =================
color_init(autoreset=True)
LOG_DIR = os.environ.get("LOG_DIR", "logs")
os.makedirs(LOG_DIR, exist_ok=True)
LOG_FILE = os.path.join(LOG_DIR, "app.log")

class ColorFormatter(logging.Formatter):
    COLORS = {
        logging.DEBUG: Fore.BLUE,
        logging.INFO: Fore.CYAN,
        logging.WARNING: Fore.YELLOW,
        logging.ERROR: Fore.RED,
        logging.CRITICAL: Fore.RED + Style.BRIGHT,
    }
    def format(self, record: logging.LogRecord) -> str:
        msg = super().format(record)
        return f"{self.COLORS.get(record.levelno, '')}{msg}{Style.RESET_ALL}"

def build_logger() -> logging.Logger:
    from logging.handlers import RotatingFileHandler
    logger = logging.getLogger("main")
    logger.setLevel(logging.DEBUG)

    # 控制台：INFO+
    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(logging.INFO)
    ch.setFormatter(ColorFormatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s", "%H:%M:%S"))

    # 文件：DEBUG+
    fh = RotatingFileHandler(LOG_FILE, maxBytes=5*1024*1024, backupCount=5, encoding="utf-8")
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s", "%Y-%m-%d %H:%M:%S"))

    # 防重复添加
    if not logger.handlers:
        logger.addHandler(ch)
        logger.addHandler(fh)
    return logger

logger = build_logger()

# ================= 配置中心 =================
@dataclasses.dataclass
class Settings:
    youtube_key: Optional[str] = None
    bitrate_k: int = 5000
    audio_bitrate_k: int = 128
    resolution: str = "720x1280"
    fps: int = 30
    gop: int = 60
    enable_recording: bool = False
    list_path: str = "list.ini"
    cookies: str = ""  # 原始字符串或 "k=v;..."
    proxy: Optional[str] = None  # eg. "http://127.0.0.1:7890"
    record_dir: str = "records"

    # 退避与探测间隔（秒）
    no_live_retry_base: float = 10.0
    no_live_retry_cap: float = 120.0

    # FFMPEG 超时/网络参数
    rw_timeout_ms: int = 10_000_000
    reconnect: bool = True

    @classmethod
    def load(cls) -> "Settings":
        cfg = configparser.ConfigParser()
        cfg.read("config.ini", encoding="utf-8")

        def g(section: str, key: str, fallback=None):
            return cfg.get(section, key, fallback=fallback) if cfg.has_option(section, key) else fallback

        s = cls(
            youtube_key=os.environ.get("YOUTUBE_KEY") or (g("youtube", "key") or None),
            bitrate_k=int(os.environ.get("VIDEO_BITRATE_K", g("youtube", "bitrate", fallback=5000))),
            audio_bitrate_k=int(os.environ.get("AUDIO_BITRATE_K", g("youtube", "audio_bitrate", fallback=128))),
            resolution=os.environ.get("VIDEO_RES", g("youtube", "resolution", fallback="720x1280")),
            fps=int(os.environ.get("FPS", g("youtube", "fps", fallback=30))),
            gop=int(os.environ.get("GOP", g("youtube", "gop", fallback=60))),
            enable_recording=(os.environ.get("ENABLE_RECORDING") or g("record", "enable", fallback="false")).lower() in {"1","true","yes","on"},
            list_path=os.environ.get("LIST_PATH", g("general", "list", fallback="list.ini")),
            cookies=os.environ.get("COOKIES", g("douyin", "cookies", fallback="")),
            proxy=os.environ.get("PROXY", g("network", "proxy", fallback=None)),
            record_dir=os.environ.get("RECORD_DIR", g("record", "dir", fallback="records")),
            no_live_retry_base=float(os.environ.get("NO_LIVE_BASE", g("backoff", "base", fallback=10.0))),
            no_live_retry_cap=float(os.environ.get("NO_LIVE_CAP", g("backoff", "cap", fallback=120.0))),
            rw_timeout_ms=int(os.environ.get("RW_TIMEOUT_MS", g("ffmpeg", "rw_timeout_ms", fallback=10_000_000))),
            reconnect=(os.environ.get("FF_RECONNECT", g("ffmpeg", "reconnect", fallback="true")).lower() in {"1","true","yes","on"}),
        )
        os.makedirs(s.record_dir, exist_ok=True)
        return s

SETTINGS = Settings.load()
if not SETTINGS.youtube_key:
    logger.error(Fore.RED + "config.ini 中无 YouTube Key，或未设置环境变量 YOUTUBE_KEY" + Style.RESET_ALL)
    sys.exit(1)

# ============== 工具函数 ==============
_USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/140.0.0.0 Safari/537.36"
)
_HEADERS_RAW = "Referer: https://www.douyin.com/\r\nOrigin: https://www.douyin.com\r\n"

async def async_sleep(sec: float):
    try:
        await asyncio.sleep(sec)
    except asyncio.CancelledError:
        pass

# ============== Douyin 解析 ==============
async def get_douyin_real_url(record_url: str, index: int, total: int,
                              record_quality: str = "原画",
                              proxy_addr: Optional[str] = None,
                              cookies: str = "") -> Tuple[Optional[str], Optional[str]]:
    """解析抖音真实流 URL 与昵称。"""
    def clip(s: str, n: int) -> str:
        return (s[:n] + "...") if len(s) > n else s

    display_url = clip(record_url, 40)
    sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [解析中...]")
    sys.stdout.flush()

    headers = {
        "User-Agent": _USER_AGENT,
        "Referer": "https://www.douyin.com/",
        "Origin": "https://www.douyin.com",
    }

    session_args = {"headers": headers}
    if proxy_addr:
        session_args["proxy"] = proxy_addr

    async with aiohttp.ClientSession(**session_args) as session:
        try:
            async with session.get(record_url, cookies=cookies, timeout=10) as resp:
                text = await resp.text()
                if not text.strip() or "captcha" in text.lower():
                    sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.RED}失败: 验证码{Style.RESET_ALL}]")
                    sys.stdout.flush()
                    return None, None

            # 选择解析函数（与原工程保持一致）
            fn = spider.get_douyin_stream_data if ('v.douyin.com' not in record_url and '/user/' not in record_url) else spider.get_douyin_app_stream_data
            json_data = await fn(url=record_url, proxy_addr=proxy_addr, cookies=cookies)
            if not json_data or not isinstance(json_data, dict):
                sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.RED}失败: 无流{Style.RESET_ALL}]")
                sys.stdout.flush()
                return None, None

            nickname = (
                json_data.get('room_title') or
                json_data.get('title') or
                json_data.get('user', {}).get('nickname') or
                json_data.get('owner', {}).get('nickname') or
                json_data.get('stream', {}).get('owner', {}).get('nickname') or
                "未知主播"
            )
            sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.YELLOW}{nickname}{Style.RESET_ALL}] [拉流...]")
            sys.stdout.flush()

            port_info = await stream.get_douyin_stream_url(json_data, record_quality, proxy_addr)
            if port_info and isinstance(port_info, dict) and port_info.get("record_url"):
                sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.GREEN}{nickname}{Style.RESET_ALL}] [成功]\n")
                sys.stdout.flush()
                return port_info.get("record_url"), nickname

            sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.RED}{nickname}{Style.RESET_ALL}] [无效流]\n")
            sys.stdout.flush()
            return None, None
        except Exception as e:
            logger.debug(f"get_douyin_real_url error: {e!r}")
            sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{Fore.RED}错误{Style.RESET_ALL}]\n")
            sys.stdout.flush()
            return None, None

# ============== 推流器 ==============
class AsyncStreamer:
    def __init__(self,
                 real_url: str,
                 youtube_key: str,
                 
                 room_name: Optional[str] = None,
                 bitrate_k: int = SETTINGS.bitrate_k,
                 audio_bitrate_k: int = SETTINGS.audio_bitrate_k,
                 resolution: str = SETTINGS.resolution,
                 fps: int = SETTINGS.fps,
                 gop: int = SETTINGS.gop,
                 enable_recording: bool = SETTINGS.enable_recording,
                 ) -> None:
        self.real_url = real_url
        self.youtube_key = youtube_key
        self.room = room_name or f"抖音_{str(uuid.uuid4())[:8]}"
        self.bitrate_k = bitrate_k
        self.audio_bitrate_k = audio_bitrate_k
        self.resolution = resolution
        self.fps = fps
        self.gop = gop
        self.enable_recording = enable_recording

        self.push_proc: Optional[asyncio.subprocess.Process] = None
        self.rec_proc: Optional[asyncio.subprocess.Process] = None
        self.pushing: bool = False
        self.start_ts: float = 0.0
        self._monitor_task: Optional[asyncio.Task] = None
        self._stderr_tasks: List[asyncio.Task] = []
        self._stop = asyncio.Event()
        self.logger = logging.getLogger(f"Streamer[{self.room}]")

    # ---------- 构建命令 ----------
    def _push_cmd(self) -> List[str]:
        flags = [
            "ffmpeg", "-y",
            "-rw_timeout", str(SETTINGS.rw_timeout_ms),
            "-user_agent", _USER_AGENT,
            "-headers", _HEADERS_RAW,
            "-protocol_whitelist", "file,pipe,http,https,tcp,tls,udp,rtp,crypto",
            "-fflags", "+discardcorrupt+genpts",
            "-flags", "low_delay",
        ]
        if SETTINGS.reconnect:
            flags += [
                "-reconnect", "1",
                "-reconnect_streamed", "1",
                "-reconnect_on_network_error", "1",
                "-reconnect_delay_max", "10",
            ]
        flags += [
            "-i", self.real_url,
            "-c:v", "libx264", "-preset", "veryfast", "-tune", "zerolatency",
            "-b:v", f"{self.bitrate_k}k",
            "-maxrate", f"{self.bitrate_k}k",
            "-bufsize", f"{self.bitrate_k*2}k",
            "-r", str(self.fps),
            "-g", str(self.gop),
            "-s", self.resolution,
            "-c:a", "aac", "-b:a", f"{self.audio_bitrate_k}k", "-ar", "44100", "-ac", "2",
            "-f", "flv", f"rtmp://a.rtmp.youtube.com/live2/{self.youtube_key}",
        ]
        return flags

    def _rec_cmd(self) -> Tuple[List[str], str]:
        base = f"{self.room}_{int(time.time())}.mp4"
        out = os.path.join(SETTINGS.record_dir, base).replace("\\", "/")
        cmd = [
            "ffmpeg", "-y",
            "-rw_timeout", str(SETTINGS.rw_timeout_ms),
            "-user_agent", _USER_AGENT,
            "-headers", _HEADERS_RAW,
            "-protocol_whitelist", "file,pipe,http,https,tcp,tls,udp,rtp,crypto",
            "-fflags", "+discardcorrupt+genpts",
            "-flags", "low_delay",
            "-i", self.real_url,
            "-c:v", "libx264", "-preset", "veryfast", "-tune", "zerolatency",
            "-c:a", "aac", "-b:a", f"{self.audio_bitrate_k}k", "-ar", "44100", "-ac", "2",
            "-movflags", "+faststart",
            out,
        ]
        return cmd, out

    # ---------- 进程与监控 ----------
    async def _spawn(self, cmd: List[str]) -> asyncio.subprocess.Process:
        # Windows 上隐藏新控制台窗口
        creationflags = 0
        if os.name == 'nt':
            creationflags = 0x08000000  # CREATE_NO_WINDOW
        return await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.PIPE,
            creationflags=creationflags,
        )

    async def _read_stderr(self, proc: asyncio.subprocess.Process, kind: str):
        assert proc.stderr is not None
        while not self._stop.is_set():
            line = await proc.stderr.readline()
            if not line:
                break
            txt = line.decode(errors='ignore').strip()
            if kind == "push" and (not self.pushing) and "frame=" in txt:
                self.pushing = True
                self.logger.info(Fore.GREEN + "成功推流到 YouTube" + Style.RESET_ALL)
            self.logger.debug(f"{kind}> {txt}")

    async def _monitor(self):
        while not self._stop.is_set():
            uptime = int(time.time() - self.start_ts) if self.start_ts else 0
            m, s = divmod(uptime, 60)
            status = f"{Fore.GREEN}推流中{Style.RESET_ALL}" if self.pushing else f"{Fore.YELLOW}启动中{Style.RESET_ALL}"
            sys.stdout.write(f"\r[{self.room}] {status} | 时间: {m:02d}:{s:02d} | 码率: {self.bitrate_k}k")
            sys.stdout.flush()
            await async_sleep(1)

    async def run(self):
        self.logger.info(Fore.CYAN + f"启动推流到 YouTube" + Style.RESET_ALL)
        self.start_ts = time.time()
        self.pushing = False

        # 启动推流
        self.push_proc = await self._spawn(self._push_cmd())
        self.logger.info(Fore.CYAN + f"推流进程 PID: {self.push_proc.pid}" + Style.RESET_ALL)
        self._stderr_tasks.append(asyncio.create_task(self._read_stderr(self.push_proc, "push")))

        # 启动录制（可选）
        if self.enable_recording:
            rec_cmd, out = self._rec_cmd()
            self.rec_proc = await self._spawn(rec_cmd)
            self.logger.info(Fore.CYAN + f"录制到: {out}" + Style.RESET_ALL)
            self._stderr_tasks.append(asyncio.create_task(self._read_stderr(self.rec_proc, "record")))

        # 监控任务
        self._monitor_task = asyncio.create_task(self._monitor())

        # 进程生命周期
        try:
            while not self._stop.is_set():
                # 推流意外退出 → 结束
                if self.push_proc and (self.push_proc.returncode is not None):
                    self.logger.warning(Fore.YELLOW + f"推流进程退出: {self.push_proc.returncode}" + Style.RESET_ALL)
                    await self.stop()
                    break
                # 录制退出 → 重启录制
                if self.enable_recording and self.rec_proc and (self.rec_proc.returncode is not None):
                    self.logger.warning(Fore.YELLOW + "录制进程退出，重启中..." + Style.RESET_ALL)
                    rec_cmd, _ = self._rec_cmd()
                    self.rec_proc = await self._spawn(rec_cmd)
                    self._stderr_tasks.append(asyncio.create_task(self._read_stderr(self.rec_proc, "record")))
                await async_sleep(1)
        finally:
            await self.stop()

    async def stop(self):
        if self._stop.is_set():
            return
        self._stop.set()
        for t in self._stderr_tasks:
            t.cancel()
        if self._monitor_task:
            self._monitor_task.cancel()

        async def _terminate(proc: Optional[asyncio.subprocess.Process]):
            if not proc:
                return
            if proc.returncode is None:
                with contextlib.suppress(ProcessLookupError):
                    proc.terminate()
                try:
                    await asyncio.wait_for(proc.wait(), timeout=2)
                except asyncio.TimeoutError:
                    with contextlib.suppress(ProcessLookupError):
                        proc.kill()
                    with contextlib.suppress(asyncio.TimeoutError):
                        await asyncio.wait_for(proc.wait(), timeout=2)
        await _terminate(self.push_proc)
        await _terminate(self.rec_proc)
        self.push_proc = None
        self.rec_proc = None
        # 清理一行状态输出
        sys.stdout.write("\r" + " " * 100 + "\r")
        sys.stdout.flush()
        self.logger.info(Fore.YELLOW + "已停止推流/录制" + Style.RESET_ALL)

# ============== 主循环 ==============
async def load_list_if_changed(path: str, last_mtime: float) -> Tuple[Optional[List[str]], float]:
    try:
        m = os.path.getmtime(path)
        if m <= last_mtime:
            return None, last_mtime
        with open(path, "r", encoding="utf-8") as f:
            urls = [ln.strip() for ln in f if ln.strip() and not ln.strip().startswith('#')]
        logger.info(Fore.CYAN + f"加载 {len(urls)} 个唯一 URL" + Style.RESET_ALL)
        return urls, m
    except FileNotFoundError:
        logger.error(Fore.RED + f"未找到 {path} 文件" + Style.RESET_ALL)
        return [], last_mtime

async def iter_until_live(urls: List[str]) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    for i, url in enumerate(urls, 1):
        real_url, nickname = await get_douyin_real_url(url, i, len(urls), proxy_addr=SETTINGS.proxy, cookies=SETTINGS.cookies)
        # 清一行
        sys.stdout.write("\r" + " " * 100 + "\r")
        sys.stdout.flush()
        if real_url and nickname:
            logger.info(Fore.GREEN + f"发现直播: {url[:50]}... [{nickname}]" + Style.RESET_ALL)
            return url, real_url, nickname
    return None, None, None

async def main():
    # 信号处理
    stop_evt = asyncio.Event()
    def _sig_handler(*_):
        stop_evt.set()
    for sig in (signal.SIGINT, signal.SIGTERM):
        with contextlib.suppress(NotImplementedError):
            asyncio.get_running_loop().add_signal_handler(sig, _sig_handler)

    logger.info(Fore.GREEN + "启动抖音直播检测" + Style.RESET_ALL)

    last_mtime = 0.0
    urls: List[str] = []
    streamer: Optional[AsyncStreamer] = None

    backoff = SETTINGS.no_live_retry_base

    while not stop_evt.is_set():
        urls_new, new_mtime = await load_list_if_changed(SETTINGS.list_path, last_mtime)
        if urls_new is not None:
            urls = urls_new
            last_mtime = new_mtime

        # 若没有在推流或已停止 → 尝试寻找新的直播
        if (streamer is None) or streamer._stop.is_set():
            if streamer and streamer._stop.is_set():
                logger.info(Fore.YELLOW + "直播中断，准备重新遍历 URL 列表" + Style.RESET_ALL)
                await async_sleep(5)
                streamer = None

            if not urls:
                sys.stdout.write(f"\r{Fore.YELLOW}无URL列表，{int(backoff)}秒后重试...{Style.RESET_ALL}")
                sys.stdout.flush()
                await async_sleep(backoff)
                backoff = min(backoff * 1.7 + (0.5 if backoff < 30 else 1.0), SETTINGS.no_live_retry_cap)
                continue

            raw_url, real_url, nickname = await iter_until_live(urls)
            if not real_url:
                sys.stdout.write(f"\r{Fore.YELLOW}未找到直播，{int(backoff)}秒后重试...{Style.RESET_ALL}")
                sys.stdout.flush()
                await async_sleep(backoff)
                backoff = min(backoff * 1.7 + (0.5 if backoff < 30 else 1.0), SETTINGS.no_live_retry_cap)
                continue

            backoff = SETTINGS.no_live_retry_base  # 命中直播 → 重置退避
            logger.info(Fore.GREEN + f"推流自: {real_url[:50]}... [{nickname}]" + Style.RESET_ALL)
            streamer = AsyncStreamer(
                real_url=real_url,
                youtube_key=SETTINGS.youtube_key,
                room_name=nickname,
                bitrate_k=SETTINGS.bitrate_k,
                audio_bitrate_k=SETTINGS.audio_bitrate_k,
                resolution=SETTINGS.resolution,
                fps=SETTINGS.fps,
                gop=SETTINGS.gop,
                enable_recording=SETTINGS.enable_recording,
            )
            asyncio.create_task(streamer.run())
        else:
            # 正在推流 → 间隔检查
            await async_sleep(10)

    # 退出
    if streamer:
        await streamer.stop()
    logger.info(Fore.YELLOW + "进程已退出" + Style.RESET_ALL)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
