import yt_dlp
import os
import re
import sys
import json
import time
import queue
import threading
import subprocess
from urllib.parse import urlparse
from flask import Flask, render_template, request, jsonify, Response, stream_with_context, send_file

app = Flask(__name__)

DEFAULT_OUTPUT_DIR = r"D:\UserData\Desktop\xiazai"
TWITTER_COOKIE = "x_cookies.txt"
BILIBILI_COOKIE = "bilibili_cookies.txt"
PROXY_URL = "socks5://127.0.0.1:10808"


def check_ffmpeg():
    """检查 FFmpeg 是否可用"""
    try:
        result = subprocess.run(['ffmpeg', '-version'], capture_output=True, text=True, timeout=5)
        return result.returncode == 0
    except Exception:
        return False


def is_twitter_url(url):
    """检查是否是 Twitter/X 链接"""
    domain = urlparse(url).netloc.lower()
    return 'x.com' in domain or 'twitter.com' in domain


def is_bilibili_url(url):
    """检查是否是 Bilibili 链接"""
    domain = urlparse(url).netloc.lower()
    return 'bilibili.com' in domain or 'b23.tv' in domain


def _find_files_by_title(output_dir, safe_title):
    """根据安全标题查找目录中的相关文件"""
    return [
        f for f in os.listdir(output_dir)
        if f.startswith(safe_title) and os.path.isfile(os.path.join(output_dir, f))
    ]


def _merge_mp4_m4a(output_dir, safe_title):
    """
    查找同名的 .mp4 和 .m4a 并用 FFmpeg 合并。
    返回 (merged: bool, message: str, filepath: str)
    """
    files = _find_files_by_title(output_dir, safe_title)
    mp4_files = [f for f in files if f.endswith('.mp4')]
    m4a_files = [f for f in files if f.endswith('.m4a')]

    if not mp4_files or not m4a_files:
        return False, "未检测到需要合并的 .mp4 + .m4a"

    # 取最可能匹配的一对：文件名最短（不含 quality suffix 的通常最短）
    mp4_files.sort(key=len)
    m4a_files.sort(key=len)
    mp4_path = os.path.join(output_dir, mp4_files[0])
    m4a_path = os.path.join(output_dir, m4a_files[0])

    # 合并输出文件名
    merged_name = safe_title + "_merged.mp4"
    merged_path = os.path.join(output_dir, merged_name)
    counter = 1
    while os.path.exists(merged_path):
        merged_name = f"{safe_title}_merged_{counter}.mp4"
        merged_path = os.path.join(output_dir, merged_name)
        counter += 1

    cmd = [
        'ffmpeg', '-y',
        '-i', mp4_path,
        '-i', m4a_path,
        '-c', 'copy',
        merged_path
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0 and os.path.exists(merged_path):
            # 合并成功后删除原始分片
            try:
                os.remove(mp4_path)
                os.remove(m4a_path)
            except Exception:
                pass
            return True, f"已合并为 {merged_name}", merged_path
        else:
            return False, f"FFmpeg 合并失败: {result.stderr[:200]}", ""
    except Exception as e:
        return False, f"合并异常: {str(e)}", ""


def download_video(video_url, output_dir, platform):
    """
    通用下载函数，返回 (success: bool, message: str, title: str, filepath: str)
    """
    os.makedirs(output_dir, exist_ok=True)

    ydl_opts = {
        'outtmpl': os.path.join(output_dir, '%(title)s.%(ext)s'),
        'ignoreerrors': True,
    }

    if platform == 'twitter':
        cookie_file = TWITTER_COOKIE
        if PROXY_URL:
            ydl_opts['proxy'] = PROXY_URL
    elif platform == 'bilibili':
        cookie_file = BILIBILI_COOKIE
        # Bilibili 不使用代理，且需要合并配置
        ydl_opts['merge_output_format'] = 'mp4'
        ydl_opts['postprocessors'] = [{
            'key': 'FFmpegVideoConvertor',
            'preferedformat': 'mp4',
        }]
    else:
        return False, "不支持的平台", "", ""

    if os.path.exists(cookie_file):
        ydl_opts['cookiefile'] = cookie_file
    else:
        return False, f"Cookie 文件 '{cookie_file}' 未找到", "", ""

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(video_url, download=True)
            if not info:
                return False, "未能获取视频信息，下载可能失败", "", ""

            title = info.get('title', 'Unknown Title')
            safe_title = re.sub(r'[<>:\"/\\|?*]', '_', title)
            files = _find_files_by_title(output_dir, safe_title)

            if not files:
                return False, "下载完成但未找到输出文件", title, ""

            # 确定最终视频文件路径
            # 优先取合并后的文件，否则取最大的 mp4 文件
            filepath = ""
            candidate_files = [f for f in files if f.endswith('.mp4')]
            if candidate_files:
                # 优先 merged 文件
                merged_candidates = [f for f in candidate_files if '_merged' in f]
                if merged_candidates:
                    filepath = os.path.join(output_dir, merged_candidates[0])
                else:
                    # 取最新/最大的文件
                    candidate_files.sort(key=lambda f: os.path.getsize(os.path.join(output_dir, f)), reverse=True)
                    filepath = os.path.join(output_dir, candidate_files[0])
            else:
                # 没有 mp4，取任意文件
                files.sort(key=lambda f: os.path.getsize(os.path.join(output_dir, f)), reverse=True)
                filepath = os.path.join(output_dir, files[0])

            # Bilibili 后处理：如果存在分离的 mp4 + m4a，手动合并
            extra_msg = ""
            if platform == 'bilibili':
                has_mp4 = any(f.endswith('.mp4') for f in files)
                has_m4a = any(f.endswith('.m4a') for f in files)
                if has_mp4 and has_m4a:
                    merged, merge_msg, merged_path = _merge_mp4_m4a(output_dir, safe_title)
                    extra_msg = f" ({merge_msg})" if merge_msg else ""
                    if merged and merged_path:
                        filepath = merged_path

            return True, f"下载成功！保存到 {output_dir}{extra_msg}", title, filepath

    except Exception as e:
        return False, f"下载失败: {str(e)}", "", ""


@app.route("/file")
def serve_file():
    """提供视频文件访问和下载"""
    filepath = request.args.get("path", "")
    as_download = request.args.get("download", "0") == "1"

    if not filepath:
        return "Missing path", 400

    # 安全检查：确保是绝对路径且文件存在
    filepath = os.path.abspath(filepath)
    if not os.path.exists(filepath) or not os.path.isfile(filepath):
        return "File not found", 404

    # 只允许访问视频文件（防止任意文件读取）
    ext = os.path.splitext(filepath)[1].lower()
    if ext not in ('.mp4', '.m4a', '.webm', '.mkv', '.mov', '.avi'):
        return "File type not allowed", 403

    return send_file(filepath, as_attachment=as_download)


@app.route("/")
def index():
    ffmpeg_ok = check_ffmpeg()
    return render_template("index.html", ffmpeg_ok=ffmpeg_ok)


@app.route("/download", methods=["POST"])
def download():
    data = request.get_json()
    url = data.get("url", "").strip() if data else ""
    output_dir = data.get("output_dir", DEFAULT_OUTPUT_DIR).strip() if data else DEFAULT_OUTPUT_DIR

    if not url:
        return jsonify({"success": False, "message": "请输入视频链接"}), 400

    if is_twitter_url(url):
        platform = 'twitter'
    elif is_bilibili_url(url):
        platform = 'bilibili'
    else:
        return jsonify({"success": False, "message": "仅支持 Twitter/X 和 Bilibili 视频链接"}), 400

    if not output_dir:
        output_dir = DEFAULT_OUTPUT_DIR

    success, message, title, filepath = download_video(url, output_dir, platform)
    return jsonify({"success": success, "message": message, "title": title, "filepath": filepath})


if __name__ == "__main__":
    print("=" * 50)
    print("Twitter / Bilibili 视频下载器")
    print("=" * 50)
    if check_ffmpeg():
        print("FFmpeg 已检测到")
    else:
        print("警告: 未检测到 FFmpeg")
    print(f"默认保存目录: {DEFAULT_OUTPUT_DIR}")
    print("访问 http://127.0.0.1:5000 打开下载页面")
    print("=" * 50)
    app.run(host="0.0.0.0", port=5000, debug=False, threaded=True)