# -*- encoding: utf-8 -*-

"""
多路并行版：
- 同时启动 N 路（N = youtube.ini 的 keys 数量）
- 每路独立 ffmpeg 进程、watchdog、定期刷新抖音真实播放地址
- 共享 URL 池，避免重复占用同一个主播
"""

import asyncio
import os
import sys
import signal
import time
import datetime
import re
import uuid
import threading
import subprocess
from typing import Any, List, Tuple, Optional, Dict, Set

import configparser
import urllib.request

from src import spider, stream
from src.utils import logger, Color, check_md5
from ffmpeg_install import check_ffmpeg, ffmpeg_path, current_env_path

# --------------------- 基础环境 ---------------------
version = "relay-parallel-2.0.0"
os.environ['PATH'] = ffmpeg_path + os.pathsep + current_env_path
script_path = os.path.split(os.path.realpath(sys.argv[0]))[0]
config_file = f'{script_path}/config/config.ini'
url_config_file = f'{script_path}/config/URL_config.ini'
yt_ini_file = f"{script_path}/config/youtube.ini"
text_encoding = 'utf-8-sig'
color = Color()

# ---------- YouTube 推流配置（读取 youtube.ini） ----------
_yt_cfg = configparser.ConfigParser()
with open(yt_ini_file, "r", encoding="utf-8-sig") as f:
    _yt_cfg.read_file(f)

def _yt_get(section: str, key: str, fallback=None, cast=str):
    try:
        if cast is bool:
            return _yt_cfg.getboolean(section, key, fallback=fallback)
        if cast is int:
            return _yt_cfg.getint(section, key, fallback=fallback)
        return _yt_cfg.get(section, key, fallback=fallback)
    except Exception:
        return fallback

def _yt_get_list(section: str, key: str) -> List[str]:
    raw = _yt_get(section, key, fallback="", cast=str) or ""
    # 支持逗号、空白、换行分隔
    parts = re.split(r"[,;\s]+", raw.strip())
    return [p for p in parts if p]

def _safe_int(v, lo, hi, default):
    try:
        v = int(v)
        if v < lo or v > hi:
            return default
        return v
    except Exception:
        return default

def _round_res(w: int, h: int) -> Tuple[int, int]:
    # 强制偶数，避免 1088/1918 这类尺寸（H.264 对齐问题）
    return (w & ~1, h & ~1)

def _build_youtube_rtmp_url(key: str, secure: bool, ingest: str | None):
    node = (ingest or "a").strip().lower()
    node = node if node.isalpha() and len(node) == 1 else "a"
    scheme = "rtmps" if secure else "rtmp"
    host   = f"{node}.rtmps.youtube.com" if secure else f"{node}.rtmp.youtube.com"
    return f"{scheme}://{host}/live2/{key}"

# --------------------- 简化配置读取 ---------------------
def read_config_value(cfg: configparser.RawConfigParser, section: str, option: str, default_value: Any) -> Any:
    try:
        cfg.read(config_file, encoding=text_encoding)
        if section not in cfg.sections():
            cfg.add_section(section)
        return cfg.get(section, option)
    except (configparser.NoSectionError, configparser.NoOptionError):
        cfg.set(section, option, str(default_value))
        with open(config_file, 'w', encoding=text_encoding) as f:
            cfg.write(f)
        return default_value

_CN_BOOL = {"是": True, "否": False}

cfg = configparser.RawConfigParser()
use_proxy = _CN_BOOL.get(read_config_value(cfg, '录制设置', '是否使用代理ip(是/否)', "否"), False)
proxy_addr = read_config_value(cfg, '录制设置', '代理地址', "")
proxy_addr = proxy_addr if use_proxy and proxy_addr.strip() else None
# Cookie（按需）
dy_cookie = read_config_value(cfg, 'Cookie', '抖音cookie', '')
tiktok_cookie = read_config_value(cfg, 'Cookie', 'tiktok_cookie', '')

# 可选：启动时检测系统代理（仅提示）
try:
    if not use_proxy:
        urllib.request.urlopen("https://www.google.com/", timeout=5)
        print("提示：你的网络似乎直连可用（无需代理）。")
except Exception:
    if not use_proxy:
        color.print_colored("提示：无法直连海外站，如需 TikTok 请在 config.ini 开启代理并填写代理地址。", color.YELLOW)

# --------------------- URL 列表读取 ---------------------
def parse_urls_from_file(path: str) -> List[str]:
    urls, seen = [], set()
    if not os.path.isfile(path):
        open(path, 'w', encoding=text_encoding).close()
        print("URL_config.ini 为空，请写入直播间地址后再运行。")
        return urls
    with open(path, "r", encoding=text_encoding, errors='ignore') as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#"):
                continue
            # 拆分（最多3段，取包含 https 或域名的那段）
            parts = re.split(r"[，,]\s*", s)
            cand = None
            for p in parts:
                if '://' in p or re.match(r"^[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", p):
                    cand = p
                    break
            if not cand:
                continue
            url = cand if '://' in cand else f"https://{cand}"
            if url not in seen:
                seen.add(url)
                urls.append(url)
    return urls

# --------------------- 直播解析与选源 ---------------------
def is_flv_preferred_platform(link: str) -> bool:
    return any(i in link for i in ["douyin", "tiktok"])

def select_source_url(link: str, stream_info: dict) -> Optional[str]:
    if not stream_info:
        return None
    if is_flv_preferred_platform(link):
        flv = stream_info.get("flv_url")
        if flv:
            try:
                from urllib.parse import urlparse, parse_qs
                qs = parse_qs(urlparse(flv).query or "")
                if (qs.get("codec") or [""])[0].lower() == "h265":
                    logger.warning("FLV 不支持 h265，改用 HLS。")
                else:
                    return flv
            except Exception:
                return flv
    return stream_info.get("record_url") or stream_info.get("m3u8_url") or stream_info.get("flv_url")

async def fetch_port_info(url: str, record_quality_code: str = "OD") -> dict:
    proxy = proxy_addr
    port_info = {}
    try:
        if "douyin.com/" in url:
            if 'v.douyin.com' not in url and '/user/' not in url:
                j = await spider.get_douyin_stream_data(url=url, proxy_addr=proxy, cookies=dy_cookie)
            else:
                j = await spider.get_douyin_app_stream_data(url=url, proxy_addr=proxy, cookies=dy_cookie)
            port_info = await stream.get_douyin_stream_url(j, record_quality_code, proxy)
        elif url.startswith("https://www.tiktok.com/"):
            if proxy:
                j = await spider.get_tiktok_stream_data(url=url, proxy_addr=proxy, cookies=tiktok_cookie)
                port_info = await stream.get_tiktok_stream_url(j, record_quality_code, proxy)
            else:
                logger.error("TikTok 需要代理，请在 config.ini 开启并配置代理。")
                port_info = {"anchor_name": "", "is_live": False}
        elif (".m3u8" in url) or (".flv" in url):
            port_info = {
                "anchor_name": "自定义_" + str(uuid.uuid4())[:8],
                "is_live": True,
            }
            if url.endswith(".flv"):
                port_info["flv_url"] = url
            else:
                port_info["m3u8_url"] = url
        else:
            port_info = {"anchor_name": "", "is_live": False}
    except Exception as e:
        logger.error(f"解析失败: {e}")
        port_info = {"anchor_name": "", "is_live": False}
    return port_info

# --------------------- ffmpeg 输出参数生成（可复用） ---------------------
def build_youtube_output_for_key(key: str, resolution: Optional[str]) -> Tuple[str, list, list]:
    if not key:
        logger.error("错误: youtube.ini 未配置 key/keys")
        sys.exit(1)

    secure  = _yt_get("youtube", "secure", fallback=False, cast=bool)
    ingest  = _yt_get("youtube", "ingest", fallback="a", cast=str)
    res_str = (resolution or _yt_get("youtube", "resolution", fallback="", cast=str) or "").strip()

    bv   = _safe_int(_yt_get("youtube", "bitrate", fallback=5000, cast=int), 300, 20000, 5000)
    ba   = _safe_int(_yt_get("youtube", "audio_bitrate", fallback=128, cast=int), 64, 384, 128)
    fps  = _safe_int(_yt_get("youtube", "fps", fallback=30, cast=int), 10, 60, 30)
    gop  = _safe_int(_yt_get("youtube", "gop", fallback=2*fps, cast=int), fps, 4*fps, 2*fps)

    out_url = _build_youtube_rtmp_url(key, secure, ingest)

    v_args = [
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-tune", "zerolatency",
        "-profile:v", "high",
        "-level:v", "4.2",
        "-pix_fmt", "yuv420p",
        "-r", str(fps),
        "-g", str(gop),
        "-b:v", f"{bv}k",
        "-maxrate", f"{bv}k",
        "-bufsize", f"{max(2*bv, 4000)}k",
    ]
    a_args = [
        "-c:a", "aac",
        "-b:a", f"{ba}k",
        "-ar", "44100",
        "-ac", "2",
    ]

    if res_str:
        try:
            w, h = map(int, res_str.lower().split("x"))
            w, h = _round_res(w, h)
            vf = (
                f"scale={w}:{h}:force_original_aspect_ratio=decrease,"
                f"pad={w}:{h}:(ow-iw)/2:(oh-ih)/2,"
                f"crop={w}:{h},"
                f"format=yuv420p"
            )
            v_args = ["-vf", vf] + v_args
        except Exception:
            logger.warning(f"resolution '{res_str}' 非法，忽略")

    return out_url, v_args, a_args

# --------------------- 启动 ffmpeg ---------------------
def launch_ffmpeg(input_url: str,
                  out_url: str,
                  extra_headers: Optional[str],
                  v_args: list,
                  a_args: list) -> subprocess.Popen:
    user_agent = ("Mozilla/5.0 (Linux; Android 11; SM-G973U) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/87.0.4280.141 Mobile Safari/537.36")
    base = [
        "ffmpeg", "-nostdin", "-hide_banner",
        "-stats", "-stats_period", "60",
        "-progress", "pipe:1",
        "-loglevel", "info",

        # 输入容错 & 自动重连
        "-rw_timeout", "15000000",
        "-timeout", "10000000",
        "-reconnect", "1",
        "-reconnect_streamed", "1",
        "-reconnect_at_eof", "1",
        "-reconnect_delay_max", "30",

        "-user_agent", user_agent,
        "-protocol_whitelist", "file,http,https,tcp,tls,crypto,udp,rtmp,rtp,httpproxy",

        "-thread_queue_size", "4096",
        "-analyzeduration", "10000000",
        "-probesize", "10000000",
    ]
    if extra_headers:
        base += ["-headers", extra_headers]

    env = os.environ.copy()
    if proxy_addr:
        env["http_proxy"] = proxy_addr
        env["https_proxy"] = proxy_addr
        base = ["ffmpeg", "-http_proxy", proxy_addr] + base[1:]

    base += ["-re", "-i", input_url]

    base += [
        "-fflags", "+genpts+igndts+discardcorrupt+nobuffer",
        "-xerror",
        "-use_wallclock_as_timestamps", "1",
        "-avioflags", "direct",
        "-rtbufsize", "100M",
    ]

    base += v_args + a_args
    base += ["-f", "flv", out_url]

    cmd_str = " ".join(base)
    print("ffmpeg cmd:", cmd_str if len(cmd_str) < 500 else cmd_str[:500] + " ...")
    logger.info(f"启动推流：{input_url} -> {out_url}")

    creationflags = 0
    if os.name == "nt":
        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP

    p = subprocess.Popen(
        base,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=0,
        env=env,
        creationflags=creationflags
    )
    p.last_out_ts = time.time()

    def _reader():
        try:
            for line in iter(p.stdout.readline, b''):
                if not line:
                    break
                text = line.decode(errors="ignore").strip()
                p.last_out_ts = time.time()

                if ("bitrate=" in text) or ("frame=" in text) or ("speed=" in text):
                    logger.info(f"[ffmpeg统计] {text}")
                if "Server disconnected" in text:
                    logger.error(f"[输出异常] {text}")
                elif "not supported" in text or "Invalid" in text or "invalid" in text:
                    logger.error(f"[转码异常] {text}")
        except Exception:
            pass

    threading.Thread(target=_reader, daemon=True).start()
    return p

def stop_ffmpeg(p: Optional[subprocess.Popen]):
    if not p:
        return
    try:
        if os.name == "nt":
            if p.stdin:
                try: p.stdin.write(b"q"); p.stdin.flush()
                except Exception: pass
            p.wait(timeout=6)
        else:
            p.send_signal(signal.SIGINT)
            p.wait(timeout=10)
    except Exception:
        try: p.kill()
        except Exception: pass

# --------------------- 平台头部 ---------------------
def headers_for_platform(platform_name: str, live_url: str) -> Optional[str]:
    live_domain = '/'.join(live_url.split('/')[0:3])
    record_headers = {
        'PandaTV': 'origin:https://www.pandalive.co.kr',
        'WinkTV': 'origin:https://www.winktv.co.kr',
        'PopkonTV': 'origin:https://www.popkontv.com',
        'FlexTV': 'origin:https://www.flextv.co.kr',
        '千度热播': 'referer:https://qiandurebo.com',
        '17Live': 'referer:https://17.live/en/live/6302408',
        '浪Live': 'referer:https://www.lang.live',
        'shopee': f'origin:{live_domain}',
        'Blued直播': 'referer:https://app.blued.cn'
    }
    return record_headers.get(platform_name)

def detect_platform(url: str) -> str:
    if "douyin.com" in url: return "抖音直播"
    if "tiktok.com" in url: return "TikTok直播"
    if url.endswith(".flv") or url.endswith(".m3u8"): return "自定义录制直播"
    return "未知平台"

# --------------------- 并行协调器 ---------------------
class UrlCoordinator:
    """ 管理 URL 池、热更新、分配给不同 worker，避免重复占用 """
    def __init__(self, path: str):
        self.path = path
        self.urls: List[str] = parse_urls_from_file(path)
        self.idx: int = 0
        self.in_use: Set[str] = set()
        self.lock = threading.Lock()
        self.last_md5 = check_md5(path)

    def refresh_if_changed(self):
        try:
            md5 = check_md5(self.path)
            if md5 != self.last_md5:
                self.urls = parse_urls_from_file(self.path) or self.urls
                self.last_md5 = md5
                print(f"[{now_str()}] 已重载 URL 列表，数量：{len(self.urls)}")
        except Exception as e:
            logger.warning(f"URL 列表热更新失败：{e}")

    def acquire_next_live(self, quality_code: str = "OD") -> Optional[Tuple[str, dict]]:
        """ 顺序扫描，跳过 in_use，找到正在直播的返回并占用 """
        with self.lock:
            if not self.urls:
                return None
            start = self.idx
            tried = 0
            while tried < len(self.urls):
                url = self.urls[self.idx]
                self.idx = (self.idx + 1) % len(self.urls)
                tried += 1
                if url in self.in_use:
                    continue
                # 释放锁去解析，避免长时间占用
                pass_url = url
                break
            else:
                return None
        # 锁外解析
        port = asyncio.run(fetch_port_info(pass_url, quality_code))
        if not port.get("anchor_name") or not port.get("is_live"):
            return None
        with self.lock:
            # 二次检查，确保没被别人抢占
            if pass_url in self.in_use:
                return None
            self.in_use.add(pass_url)
        return pass_url, port

    def release(self, url: Optional[str]):
        if not url:
            return
        with self.lock:
            self.in_use.discard(url)

def now_str(): return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# --------------------- Worker ---------------------
class StreamWorker(threading.Thread):
    def __init__(self, worker_id: int, key: str, resolution: Optional[str], coordinator: UrlCoordinator):
        super().__init__(daemon=True)
        self.worker_id = worker_id
        self.key = key
        self.resolution = resolution
        self.coordinator = coordinator

        self.current_url: Optional[str] = None
        self.current_real_url: Optional[str] = None
        self.current_proc: Optional[subprocess.Popen] = None
        self.current_platform = ""
        self.current_anchor = ""
        self.record_quality_code = "OD"

        self.watchdog_grace = 120
        self.watchdog_silence = 600
        self.refresh_real_url_every = 300
        self.check_interval_live = 8
        self.scan_interval_when_idle = 2

        self._stop_event = threading.Event()

    def stop(self):
        self._stop_event.set()

    def run(self):
        color.print_colored(f"[Worker-{self.worker_id}] 启动", color.GREEN)
        try:
            while not self._stop_event.is_set():
                # 刷新 URL 文件
                self.coordinator.refresh_if_changed()

                if self.current_proc and self.current_url:
                    retcode = self.current_proc.poll()
                    if retcode is not None:
                        color.print_colored(f"[{now_str()}][Worker-{self.worker_id}] ffmpeg 退出(code={retcode})，释放并切源…", color.YELLOW)
                        stop_ffmpeg(self.current_proc)
                        self.coordinator.release(self.current_url)
                        self._clear_state()
                        continue

                    # 仍在推流：检查主播是否仍在直播
                    still_live = True
                    try:
                        port = asyncio.run(fetch_port_info(self.current_url, self.record_quality_code))
                        still_live = bool(port.get("is_live"))
                    except Exception as e:
                        logger.error(f"[Worker-{self.worker_id}] 检测当前源失败：{e}")
                        still_live = True

                    if still_live:
                        elapsed = int(time.time() - self.start_time)
                        h, m = divmod(elapsed, 3600)
                        m, s = divmod(m, 60)
                        dur_str = f"{h:02d}:{m:02d}:{s:02d}"
                        print(f"\r[{now_str()}][Worker-{self.worker_id}] 正在转播({dur_str})：{self.current_platform} | {self.current_anchor} | {self.current_url}", end="")

                        # 定期刷新真实播放地址
                        try:
                            if elapsed >= self.refresh_real_url_every and (elapsed % self.refresh_real_url_every) < self.check_interval_live:
                                new_port = asyncio.run(fetch_port_info(self.current_url, self.record_quality_code))
                                new_real_url = select_source_url(self.current_url, new_port)
                                if new_real_url and new_real_url != self.current_real_url:
                                    color.print_colored(f"\n[{now_str()}][Worker-{self.worker_id}] 刷新播放地址（token 变化），重启 ffmpeg…", color.YELLOW)
                                    stop_ffmpeg(self.current_proc)
                                    out_url, v_args, a_args = build_youtube_output_for_key(self.key, self.resolution)
                                    headers = headers_for_platform(self.current_platform, self.current_url)
                                    self.current_proc = launch_ffmpeg(new_real_url, out_url, headers, v_args, a_args)
                                    self.current_real_url = new_real_url
                                    self.start_time = time.time()
                                    time.sleep(self.check_interval_live)
                                    continue
                        except Exception as e:
                            logger.warning(f"[Worker-{self.worker_id}] 刷新播放地址失败：{e}")

                        # watchdog：心跳超时才重启
                        silence = time.time() - getattr(self.current_proc, "last_out_ts", 0)
                        if elapsed >= self.watchdog_grace and silence > self.watchdog_silence:
                            color.print_colored(f"\n[{now_str()}][Worker-{self.worker_id}] 超过10分钟未见输出，重启…", color.YELLOW)
                            stop_ffmpeg(self.current_proc)
                            self._clear_state(release=True)
                            time.sleep(1)
                            continue

                        time.sleep(self.check_interval_live)
                        continue
                    else:
                        color.print_colored(f"\n[{now_str()}][Worker-{self.worker_id}] 主播停播，释放并切源…", color.YELLOW)
                        stop_ffmpeg(self.current_proc)
                        self._clear_state(release=True)
                        continue

                # 未在推流：尝试分配新的直播 URL
                got = self.coordinator.acquire_next_live(self.record_quality_code)
                if not got:
                    print(f"[{now_str()}][Worker-{self.worker_id}] 无可用直播源，稍后再试…")
                    time.sleep(self.scan_interval_when_idle)
                    continue

                url, port = got
                real_url = select_source_url(url, port)
                if not real_url:
                    print(f"[{now_str()}][Worker-{self.worker_id}] 未得到播放地址，跳过：{url}")
                    self.coordinator.release(url)
                    time.sleep(1)
                    continue

                self.current_platform = detect_platform(url)
                self.current_anchor = port.get("anchor_name", "")
                headers = headers_for_platform(self.current_platform, url)
                out_url, v_args, a_args = build_youtube_output_for_key(self.key, self.resolution)
                self.current_proc = launch_ffmpeg(real_url, out_url, headers, v_args, a_args)
                self.current_url = url
                self.current_real_url = real_url
                self.start_time = time.time()
                color.print_colored(f"[{now_str()}][Worker-{self.worker_id}] 开始转播：{self.current_platform} | {self.current_anchor} | {url}", color.GREEN)

        except Exception as e:
            logger.error(f"[Worker-{self.worker_id}] 主循环异常：{e}")
        finally:
            stop_ffmpeg(self.current_proc)
            self.coordinator.release(self.current_url)
            print(f"[Worker-{self.worker_id}] 已退出。")

    def _clear_state(self, release: bool = False):
        if release and self.current_url:
            self.coordinator.release(self.current_url)
        self.current_url = None
        self.current_real_url = None
        self.current_proc = None
        self.current_platform = ""
        self.current_anchor = ""

# --------------------- 启动入口 ---------------------
def check_ffmpeg_existence() -> bool:
    try:
        result = subprocess.run(['ffmpeg', '-version'], check=True, capture_output=True, text=True)
        if result.returncode == 0:
            print(result.stdout.splitlines()[0])
    except Exception:
        pass
    finally:
        return bool(check_ffmpeg())

def main():
    print("-----------------------------------------------------")
    print("|         DouyinLiveRecorder - 多路并行转推版        |")
    print("-----------------------------------------------------")
    print(f"版本号: {version}\nGitHub: https://github.com/ihmily/DouyinLiveRecorder")
    print("模式：并行检测 + 并行推流（每路独立 watchdog）")
    print(".....................................................")

    if not check_ffmpeg_existence():
        logger.error("缺少 ffmpeg，程序退出")
        sys.exit(1)

    urls = parse_urls_from_file(url_config_file)
    if not urls:
        logger.error("URL_config.ini 中没有可用 URL。")
        sys.exit(1)

    keys = _yt_get_list("youtube", "keys")
    if not keys:
        # 兼容单 key 配置
        single_key = _yt_get("youtube", "key", fallback="", cast=str)
        if single_key:
            keys = [single_key]
    if not keys:
        logger.error("youtube.ini 没有配置 [youtube].keys 或 key")
        sys.exit(1)

    # 可选：为每个 key 指定不同分辨率（与 keys 顺序对应），否则使用全局 resolution
    per_res = _yt_get_list("youtube", "resolutions")  # 例如：1920x1080,1080x1920,1280x720
    while per_res and len(per_res) < len(keys):
        per_res.append(per_res[-1])  # 不足时沿用最后一个

    # 协调器（共享 URL 池）
    coordinator = UrlCoordinator(url_config_file)

    # 启动多路 worker
    workers: List[StreamWorker] = []
    for i, key in enumerate(keys, start=1):
        res_for_this = per_res[i-1] if per_res else None
        w = StreamWorker(worker_id=i, key=key, resolution=res_for_this, coordinator=coordinator)
        w.start()
        workers.append(w)

    # 主线程仅负责捕获 Ctrl+C 并优雅退出
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n收到中断，正在退出…")
    finally:
        for w in workers:
            w.stop()
        for w in workers:
            w.join(timeout=10)
        print("全部 worker 已退出。")

if __name__ == "__main__":
    main()