#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VideoLingo 快速安装脚本 - 使用清华大学镜像源
解决PyTorch等大型包下载缓慢的问题
"""

import subprocess
import sys
import os
from rich.console import Console
from rich.panel import Panel
from rich.progress import track
import time

console = Console()

def run_command(cmd, description=""):
    """执行命令并显示进度"""
    console.print(f"🚀 {description}", style="bold green")
    console.print(f"执行命令: {cmd}", style="dim")
    
    try:
        result = subprocess.run(cmd, shell=True, check=True, 
                              capture_output=True, text=True)
        console.print("✅ 完成", style="bold green")
        return True
    except subprocess.CalledProcessError as e:
        console.print(f"❌ 错误: {e}", style="bold red")
        console.print(f"输出: {e.stdout}", style="dim")
        console.print(f"错误: {e.stderr}", style="dim")
        return False

def main():
    console.print(Panel.fit("🎯 VideoLingo 快速安装 - 清华镜像版", style="bold blue"))
    
    # 1. 首先配置pip使用清华镜像
    console.print("📦 配置pip使用清华大学镜像源...")
    pip_config_cmd = """
mkdir -p ~/.pip
cat > ~/.pip/pip.conf << EOF
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple/
[install]
trusted-host = pypi.tuna.tsinghua.edu.cn
EOF
"""
    run_command(pip_config_cmd, "配置pip镜像源")
    
    # 2. 先停止当前的pytorch下载（如果正在进行）
    console.print("⏹️ 准备重新安装PyTorch...")
    
    # 3. 安装PyTorch及相关组件（使用清华镜像）
    pytorch_packages = [
        "torch==2.0.0+cu118",
        "torchvision==0.15.0+cu118", 
        "torchaudio==2.0.0+cu118"
    ]
    
    pytorch_cmd = f"""pip install {' '.join(pytorch_packages)} \\
        -i https://pypi.tuna.tsinghua.edu.cn/simple/ \\
        --extra-index-url https://download.pytorch.org/whl/cu118 \\
        --timeout 300"""
    
    if not run_command(pytorch_cmd, "安装PyTorch CUDA版本"):
        console.print("⚠️ PyTorch安装失败，尝试CPU版本...", style="yellow")
        cpu_cmd = "pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple/"
        run_command(cpu_cmd, "安装PyTorch CPU版本")
    
    # 4. 安装其他可能需要的依赖
    other_packages = [
        "transformers",
        "accelerate", 
        "datasets",
        "tokenizers",
        "huggingface-hub",
        "safetensors",
        "numpy",
        "opencv-python",
        "pillow",
        "matplotlib", 
        "scikit-learn",
        "pandas",
        "tqdm",
        "psutil",
        "ffmpeg-python",
        "whisper-openai",
        "openai",
        "anthropic"
    ]
    
    console.print("📦 安装其他常用AI/ML依赖包...")
    for package in track(other_packages, description="安装依赖包..."):
        cmd = f"pip install {package} -i https://pypi.tuna.tsinghua.edu.cn/simple/ --timeout 180"
        subprocess.run(cmd, shell=True, capture_output=True)
        time.sleep(0.1)  # 避免过快请求
    
    # 5. 验证安装
    console.print("🔍 验证PyTorch安装...")
    verify_cmd = """python -c "
import torch
print(f'PyTorch版本: {torch.__version__}')
print(f'CUDA可用: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA版本: {torch.version.cuda}')
    print(f'GPU数量: {torch.cuda.device_count()}')
    print(f'当前GPU: {torch.cuda.get_device_name(0)}')
"
"""
    
    run_command(verify_cmd, "验证PyTorch安装")
    
    console.print(Panel.fit("🎉 安装完成！现在可以继续运行VideoLingo了", style="bold green"))
    console.print("💡 提示: 如果还有其他包需要安装，它们将自动使用清华镜像源", style="dim")

if __name__ == "__main__":
    main()