import requests
import re
import json
import subprocess
import shlex
import os
import sys

# ================= 配置 =================

CONFIG_FILE = "upstreams.json"
OUTPUT_FILE = "channels.json"

# ================= 功能函数 =================

def get_stream_data(url):
    """
    获取媒体详细信息
    返回字典包含: width, height, bitrate, video, audio, fps, rate
    """
    data = {
        "width": 0,
        "height": 0,
        "bitrate": 0,
        "video": "",
        "audio": "",
        "fps": 0,
        "rate": 0
    }
    
    has_video = False
    has_audio = False

    try:
        # 8秒超时检测
        cmd = f'ffprobe -v quiet -print_format json -show_format -show_streams "{url}"'
        result = subprocess.run(shlex.split(cmd), capture_output=True, text=True, timeout=8)
        
        if result.returncode != 0:
            return None 

        info = json.loads(result.stdout)
        
        if "streams" not in info or not info["streams"]:
            return None 

        for stream in info["streams"]:
            codec_type = stream.get("codec_type")
            
            # --- 视频流 ---
            if codec_type == "video" and not has_video:
                has_video = True
                data["width"] = stream.get("width", 0)
                data["height"] = stream.get("height", 0)
                data["video"] = stream.get("codec_name", "unknown")
                
                r_rate = stream.get("r_frame_rate", "0/0")
                if "/" in r_rate:
                    try:
                        num, den = map(int, r_rate.split('/'))
                        if den > 0:
                            data["fps"] = round(num / den, 2)
                    except: pass
                else:
                    try:
                        data["fps"] = float(r_rate)
                    except: pass
                
                if stream.get("bit_rate"):
                    data["bitrate"] = int(stream.get("bit_rate"))

            # --- 音频流 ---
            elif codec_type == "audio" and not has_audio:
                has_audio = True
                data["audio"] = stream.get("codec_name", "unknown")
                
                sample_rate = stream.get("sample_rate")
                if sample_rate:
                    try:
                        data["rate"] = int(sample_rate)
                    except: pass
                
                if data["bitrate"] == 0 and stream.get("bit_rate"):
                    data["bitrate"] = int(stream.get("bit_rate"))

        # 有效性判定
        if not has_video and not has_audio:
            return None

        # 码率兜底
        if data["bitrate"] == 0 and "format" in info:
            fmt_bitrate = info["format"].get("bit_rate")
            if fmt_bitrate:
                try:
                    data["bitrate"] = int(fmt_bitrate)
                except: pass

    except Exception:
        return None 

    return data

def check_keywords(targets, keywords):
    """
    检查是否命中关键词 (targets: [url, name, group])
    """
    for target in targets:
        if not target: continue
        for kw in keywords:
            if kw in target:
                return True
    return False

def parse_m3u_content(content, upstream_config, seen_urls):
    """
    解析 M3U 内容
    upstream_config: 包含 url, tag, whitelist, blacklist, replace 等信息的字典
    """
    results = []
    lines = content.splitlines()
    current_meta = {}
    
    # 提取配置信息
    tag = upstream_config.get('tag', '')
    whitelist = upstream_config.get('whitelist', [])
    blacklist = upstream_config.get('blacklist', [])
    replace_rules = upstream_config.get('replace', {})

    for line in lines:
        line = line.strip()
        if not line:
            continue

        if line.startswith("#EXTINF"):
            current_meta = {}
            # 提取所有属性 key="value"
            attributes = re.findall(r'([a-zA-Z0-9-]+)="([^"]*)"', line)
            for key, value in attributes:
                current_meta[key.lower()] = value
            
            # 提取频道名称
            last_comma = line.rfind(',')
            if last_comma != -1:
                current_meta['name_raw'] = line[last_comma+1:].strip()
            else:
                current_meta['name_raw'] = ""

        elif not line.startswith("#"):
            url = line
            
            # --- 1. 预处理 ---
            name = current_meta.get('name_raw', 'Unknown')
            group = current_meta.get('group-title', '未分类')
            
            # --- 2. 替换 ---
            if replace_rules:
                for k, v in replace_rules.items():
                    url = url.replace(k, v)
            
            # --- 3. 过滤 (White/Black List) ---
            check_targets = [url, name, group]
            
            if whitelist and len(whitelist) > 0:
                # 存在白名单：必须命中才保留
                if not check_keywords(check_targets, whitelist):
                    continue 
            elif blacklist and len(blacklist) > 0:
                # 仅存在黑名单：命中则丢弃
                if check_keywords(check_targets, blacklist):
                    continue
            
            # --- 4. 去重 ---
            if url in seen_urls:
                continue
            seen_urls.add(url)
            
            # --- 5. 名称 ---
            tvg = current_meta.get('tvg-name') or current_meta.get('tvg-id') or name

            print(f"检测 [{tag}]: {name} | {url}", end="", flush=True)

            # --- 6. FFprobe 检测 ---
            s_data = get_stream_data(url)

            if s_data is None:
                print(" -> [无效] PASS")
                continue
            
            # --- 7. 输出与保存 ---
            out_info = []
            if s_data['width'] > 0:
                out_info.append(f"视频:{s_data['video']} {s_data['width']}x{s_data['height']}")
            else:
                out_info.append("视频:无")
                
            if s_data['audio']:
                out_info.append(f"音频:{s_data['audio']}")
            
            print(f" -> [OK] {'|'.join(out_info)}")

            channel_obj = {
                "tag": tag,            # 新增: 来源标签
                "group": group,
                "tvg": tvg,
                "name": name,
                "url": url,
                "width": s_data['width'],
                "height": s_data['height'],
                "bitrate": s_data['bitrate'],
                "video": s_data['video'],
                "audio": s_data['audio'],
                "fps": s_data['fps'],
                "rate": s_data['rate']
            }
            results.append(channel_obj)
            
    return results

def load_upstreams(filepath):
    """读取配置文件"""
    if not os.path.exists(filepath):
        print(f"错误: 找不到配置文件 {filepath}")
        return []
    
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return json.load(f)
    except json.JSONDecodeError as e:
        print(f"错误: 配置文件格式不正确 (JSON解析失败) - {e}")
        return []

def main():
    # 1. 读取外部 JSON 配置
    upstreams = load_upstreams(CONFIG_FILE)
    if not upstreams:
        print("未加载到有效配置，程序退出。")
        sys.exit(1)

    all_data = []
    global_seen = set()

    # 2. 遍历处理
    for upstream in upstreams:
        url = upstream.get('url')
        if not url: continue

        print(f"\n--- 正在处理源 [{upstream.get('tag', '无标签')}]: {url} ---")
        try:
            resp = requests.get(url, timeout=10)
            resp.encoding = 'utf-8'
            
            if resp.status_code == 200:
                # 将整个 upstream 配置传进去，方便提取 tag/whitelist/blacklist 等
                data = parse_m3u_content(resp.text, upstream, global_seen)
                all_data.extend(data)
            else:
                print(f"HTTP 错误: {resp.status_code}")
                
        except Exception as e:
            print(f"下载异常: {e}")

    # 3. 保存结果
    print(f"\n处理完成，共保存 {len(all_data)} 个有效频道到 {OUTPUT_FILE}")
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        json.dump(all_data, f, ensure_ascii=False, indent=4)

if __name__ == "__main__":
    main()