import os
import subprocess
import threading
import asyncio
import uuid
import time
import signal
import sys
import logging
import aiohttp
import spider  # 您的 spider 模块
import stream  # 您的 stream 模块
import configparser
from typing import Optional, Tuple
from colorama import init, Fore, Style
from logging.handlers import RotatingFileHandler

# ----------------- 日志设置 -----------------
# 初始化 colorama
init(autoreset=True)

# 创建 logs 目录
LOG_DIR = "logs"
os.makedirs(LOG_DIR, exist_ok=True)
LOG_FILE = os.path.join(LOG_DIR, "app.log")

# 自定义彩色日志 Formatter
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):
        color = self.COLORS.get(record.levelno, "")
        message = super().format(record)
        return f"{color}{message}{Style.RESET_ALL}"

# 控制台日志（INFO 以上，带颜色）
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
console_formatter = ColorFormatter("%(asctime)s [%(levelname)s] %(message)s", "%H:%M:%S")
console_handler.setFormatter(console_formatter)

# 文件日志（DEBUG 以上，自动滚动，5MB*5份）
file_handler = RotatingFileHandler(LOG_FILE, maxBytes=5*1024*1024, backupCount=5, encoding="utf-8")
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s", "%Y-%m-%d %H:%M:%S")
file_handler.setFormatter(file_formatter)

# 根 logger
logging.basicConfig(level=logging.DEBUG, handlers=[console_handler, file_handler])
logger = logging.getLogger("main")

# 抑制不必要的三方日志
logging.getLogger("aiohttp").setLevel(logging.CRITICAL)
logging.getLogger("aiohttp.client").setLevel(logging.CRITICAL)
logging.getLogger("httpcore").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)
logging.getLogger("asyncio").setLevel(logging.CRITICAL)
if hasattr(spider, 'logger'):
    spider.logger.setLevel(logging.CRITICAL)
if hasattr(stream, 'logger'):
    stream.logger.setLevel(logging.CRITICAL)

# ----------------- 主逻辑 -----------------
class DouyinToYouTubeStreamer:
    def __init__(self, bitrate=5000, audio_bitrate=128, resolution='720x1280', fps=30, gop=60, enable_recording=False):
        self.push_proc = None
        self.record_proc = None
        self.stop_flag = threading.Event()
        self.bitrate = bitrate
        self.audio_bitrate = audio_bitrate
        self.resolution = resolution
        self.fps = fps
        self.gop = gop
        self.start_time = None
        self.room_name = None
        self.pushing = False
        self.out_file = None
        self.enable_recording = enable_recording
        os.makedirs("records", exist_ok=True)
        self.logger = logging.getLogger(self.__class__.__name__)

    def start_stream(self, real_url, youtube_key, room_name=None):
        self.room_name = room_name or f"抖音_{str(uuid.uuid4())[:8]}"
        self.logger.info(f"{Fore.GREEN}[{self.room_name}] 启动推流到 YouTube{Style.RESET_ALL}")
        threading.Thread(target=self._run_stream_loop, args=(real_url, youtube_key), daemon=True).start()
        threading.Thread(target=self._monitor_thread, daemon=True).start()

    def _build_push_cmd(self, real_url, youtube_key):
        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 = "Referer: https://www.douyin.com/\r\nOrigin: https://www.douyin.com\r\n"
        return [
            'ffmpeg', '-y',
            '-rw_timeout', '10000000',
            '-user_agent', user_agent,
            '-headers', headers,
            '-protocol_whitelist', 'file,pipe,http,https,tcp,tls,udp,rtp,crypto',
            '-fflags', '+discardcorrupt+genpts',
            '-flags', 'low_delay',
            '-i', real_url,
            '-c:v', 'libx264', '-preset', 'veryfast', '-tune', 'zerolatency',
            '-b:v', f'{self.bitrate}k',
            '-maxrate', f'{self.bitrate}k',
            '-bufsize', f'{self.bitrate * 2}k',
            '-r', str(self.fps),
            '-g', str(self.gop),
            '-s', self.resolution,
            '-c:a', 'aac', '-b:a', f'{self.audio_bitrate}k', '-ar', '44100', '-ac', '2',
            '-f', 'flv',
            f'rtmp://a.rtmp.youtube.com/live2/{youtube_key}'
        ]

    def _build_record_cmd(self, real_url):
        base_name = f"{self.room_name}_{int(time.time())}.mp4"
        self.out_file = os.path.join("records", base_name).replace("\\", "/")
        return [
            'ffmpeg', '-y',
            '-i', real_url,
            '-c:v', 'libx264', '-preset', 'veryfast', '-tune', 'zerolatency',
            '-c:a', 'aac', '-b:a', f'{self.audio_bitrate}k', '-ar', '44100', '-ac', '2',
            '-movflags', '+faststart',
            self.out_file
        ]

    def _run_stream_loop(self, real_url, youtube_key):
        while not self.stop_flag.is_set():
            try:
                self.start_time = time.time()
                self.pushing = False
                push_cmd = self._build_push_cmd(real_url, youtube_key)
                self.push_proc = subprocess.Popen(
                    push_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                    text=True, bufsize=1
                )
                self.logger.info(f"{Fore.CYAN}[{self.room_name}] 推流进程 PID: {self.push_proc.pid}{Style.RESET_ALL}")

                if self.enable_recording:
                    record_cmd = self._build_record_cmd(real_url)
                    self.record_proc = subprocess.Popen(
                        record_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                        text=True, bufsize=1
                    )
                    self.logger.info(f"{Fore.CYAN}[{self.room_name}] 录制到: {self.out_file}{Style.RESET_ALL}")

                def monitor_proc(proc, name):
                    while not self.stop_flag.is_set() and proc and proc.poll() is None:
                        line = proc.stderr.readline()
                        if line and name == "push" and not self.pushing and "frame=" in line:
                            self.pushing = True
                            self.logger.info(f"{Fore.GREEN}[{self.room_name}] 成功推流到 YouTube{Style.RESET_ALL}")

                threading.Thread(target=monitor_proc, args=(self.push_proc, "push"), daemon=True).start()
                if self.enable_recording:
                    threading.Thread(target=monitor_proc, args=(self.record_proc, "record"), daemon=True).start()

                while not self.stop_flag.is_set():
                    if self.push_proc.poll() is not None:
                        self.logger.warning(f"{Fore.YELLOW}[{self.room_name}] 推流进程退出{Style.RESET_ALL}")
                        self.stop_flag.set()
                        break
                    if self.enable_recording and self.record_proc and self.record_proc.poll() is not None:
                        self.logger.warning(f"{Fore.YELLOW}[{self.room_name}] 录制进程退出，重启中...{Style.RESET_ALL}")
                        self._terminate_procs()
                        break
                    time.sleep(1)

            except Exception as e:
                self.logger.error(f"{Fore.RED}[{self.room_name}] FFmpeg 错误: {e}{Style.RESET_ALL}")
            finally:
                self._terminate_procs()
                if not self.stop_flag.is_set():
                    time.sleep(1)

    def _terminate_procs(self):
        for proc in [self.push_proc, self.record_proc]:
            if proc and proc.poll() is None:
                proc.terminate()
                try:
                    proc.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    proc.kill()
        self.push_proc, self.record_proc = None, None

    def _monitor_thread(self):
        while not self.stop_flag.is_set():
            uptime = int(time.time() - self.start_time) if self.start_time 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_name}] {status} | 时间: {m:02d}:{s:02d} | 码率: {self.bitrate}k")
            sys.stdout.flush()
            time.sleep(1)

    def stop(self):
        self.logger.info(f"{Fore.YELLOW}[{self.room_name}] 停止推流和录制{Style.RESET_ALL}")
        self.stop_flag.set()
        self._terminate_procs()

# ----------------- 其他函数保持不变 -----------------
async def get_douyin_real_url(record_url: str, index: int, total: int, record_quality="原画", proxy_addr=None, cookies='') -> Tuple[Optional[str], Optional[str]]:
    def update_status(status: str, nickname: str = "未知主播"):
        display_url = record_url[:40] + "..." if len(record_url) > 40 else record_url
        display_nickname = nickname[:20] + "..." if len(nickname) > 20 else nickname
        sys.stdout.write(f"\r{index}/{total} 检查: {display_url} [{display_nickname}] [{status}]")
        sys.stdout.flush()

    headers = {
        "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",
        "Referer": "https://www.douyin.com/",
        "Origin": "https://www.douyin.com"
    }
    session_args = {"headers": headers}
    if proxy_addr:
        session_args["proxy"] = proxy_addr

    update_status(f"{Fore.YELLOW}检测中...{Style.RESET_ALL}")

    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():
                    update_status(f"{Fore.RED}失败: 验证码{Style.RESET_ALL}")
                    return None, None

                json_data = await (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)(url=record_url, proxy_addr=proxy_addr, cookies=cookies)
                
                if not json_data or not isinstance(json_data, dict):
                    update_status(f"{Fore.RED}失败: 无流{Style.RESET_ALL}")
                    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
                    "未知主播"
                )
                update_status(f"{Fore.YELLOW}解析中...{Style.RESET_ALL}", nickname)

                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"):
                    update_status(f"{Fore.GREEN}成功{Style.RESET_ALL}", nickname)
                    return port_info.get("record_url"), nickname
                update_status(f"{Fore.RED}失败: 无效流{Style.RESET_ALL}", nickname)
                return None, None
        except Exception:
            update_status(f"{Fore.RED}失败: 错误{Style.RESET_ALL}")
            return None, None

def start_douyin_to_youtube(record_url: str, youtube_key: str, real_url: Optional[str] = None, nickname: Optional[str] = None,
                            proxy_addr=None, cookies='', bitrate=5000, enable_recording=False):
    if real_url is None or nickname is None:
        real_url, nickname = asyncio.run(get_douyin_real_url(record_url, 0, 0, proxy_addr=proxy_addr, cookies=cookies))
    if not real_url:
        logger.error(f"{Fore.RED}无法获取 {record_url[:50]}... 的流{Style.RESET_ALL}")
        return None
    logger.info(f"{Fore.GREEN}推流自: {real_url[:50]}... [{nickname}]{Style.RESET_ALL}")
    streamer = DouyinToYouTubeStreamer(bitrate=bitrate, enable_recording=enable_recording)
    streamer.start_stream(real_url, youtube_key, room_name=nickname)
    return streamer

def load_config():
    config = configparser.ConfigParser()
    config.read("config.ini", encoding="utf-8")
    return config.get("youtube", "key", fallback=None)

def load_list(last_mtime=0):
    try:
        mtime = os.path.getmtime("list.ini")
        if mtime <= last_mtime:
            return None, last_mtime
        with open("list.ini", "r", encoding="utf-8") as f:
            urls = [line.strip() for line in f if line.strip() and not line.startswith("#")]
        logger.info(f"{Fore.CYAN}加载 {len(urls)} 个唯一 URL{Style.RESET_ALL}")
        return urls, mtime
    except FileNotFoundError:
        logger.error(f"{Fore.RED}未找到 list.ini 文件{Style.RESET_ALL}")
        return [], last_mtime

if __name__ == "__main__":
    youtube_key = load_config()
    if not youtube_key:
        logger.error(f"{Fore.RED}config.ini 中无 YouTube Key{Style.RESET_ALL}")
        sys.exit(1)

    streamer_instance = None
    last_mtime = 0
    urls = []

    def signal_handler(sig, frame):
        if streamer_instance:
            streamer_instance.stop()
        sys.stdout.write("\r" + " " * 80 + "\r")
        sys.stdout.flush()
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    logger.info(f"{Fore.GREEN}启动抖音直播检测{Style.RESET_ALL}")
    while True:
        urls_result, new_mtime = load_list(last_mtime)
        if urls_result is not None:
            urls = urls_result
            last_mtime = new_mtime

        if not streamer_instance or (hasattr(streamer_instance, 'stop_flag') and streamer_instance.stop_flag.is_set()):
            if streamer_instance:
                logger.info(f"{Fore.YELLOW}直播中断，重新遍历 URL 列表{Style.RESET_ALL}")
                sys.stdout.write("\r" + " " * 80 + "\r")
                sys.stdout.flush()
                time.sleep(5)  # 中断后短暂延迟
            streamer_instance = None
            if not urls:
                sys.stdout.write(f"\r{Fore.YELLOW}无URL列表，10秒后重试...{Style.RESET_ALL}")
                sys.stdout.flush()
                time.sleep(10)
                continue

            for i, url in enumerate(urls, 1):
                real_url, nickname = asyncio.run(get_douyin_real_url(url, i, len(urls)))
                sys.stdout.write("\r" + " " * 80 + "\r")
                sys.stdout.flush()
                if real_url and nickname:
                    logger.info(f"{Fore.GREEN}发现直播: {url[:50]}... [{nickname}]{Style.RESET_ALL}")
                    streamer_instance = start_douyin_to_youtube(url, youtube_key, real_url=real_url, nickname=nickname)
                    break

            if not streamer_instance:
                sys.stdout.write(f"\r{Fore.YELLOW}未找到直播，60秒后重试...{Style.RESET_ALL}")
                sys.stdout.flush()
                time.sleep(60)
        else:
            time.sleep(10)
