#!/usr/bin/env python3
"""
环境管理脚本 - 创建和查看Python/Conda环境
"""
import os
import sys
import subprocess
import json
from pathlib import Path
from typing import Optional, List

class Colors:
    """终端颜色"""
    RED = '\033[0;31m'
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[0;34m'
    CYAN = '\033[0;36m'
    NC = '\033[0m'  # No Color

def print_colored(text: str, color: str = Colors.NC):
    """打印彩色文本"""
    print(f"{color}{text}{Colors.NC}")

def run_command(cmd: List[str], check: bool = True) -> tuple[str, int]:
    """运行命令并返回输出和返回码"""
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=check
        )
        return result.stdout.strip(), result.returncode
    except subprocess.CalledProcessError as e:
        return e.stdout.strip() + "\n" + e.stderr.strip(), e.returncode
    except FileNotFoundError:
        return "", 1

def check_command_exists(cmd: str) -> bool:
    """检查命令是否存在"""
    _, code = run_command(["which", cmd], check=False)
    return code == 0

def create_conda_env(env_name: str, python_version: str = "3.10"):
    """创建conda环境"""
    if not env_name:
        print_colored("错误: 请提供环境名称", Colors.RED)
        return False
    
    if not check_command_exists("conda"):
        print_colored("错误: 未找到conda命令，请先安装Anaconda或Miniconda", Colors.RED)
        return False
    
    print_colored(f"创建conda环境: {env_name} (Python {python_version})", Colors.GREEN)
    
    output, code = run_command([
        "conda", "create", "-n", env_name, 
        f"python={python_version}", "-y"
    ], check=False)
    
    if code == 0:
        print_colored("✓ Conda环境创建成功", Colors.GREEN)
        print()
        print_colored("激活环境:", Colors.YELLOW)
        print(f"  conda activate {env_name}")
        return True
    else:
        print_colored(f"错误: {output}", Colors.RED)
        return False

def create_venv_env(env_name: str):
    """创建venv环境"""
    if not env_name:
        print_colored("错误: 请提供环境名称", Colors.RED)
        return False
    
    print_colored(f"创建venv环境: {env_name}", Colors.GREEN)
    
    try:
        subprocess.run([sys.executable, "-m", "venv", env_name], check=True)
        print_colored("✓ Venv环境创建成功", Colors.GREEN)
        print()
        print_colored("激活环境:", Colors.YELLOW)
        print(f"  source {env_name}/bin/activate")
        return True
    except subprocess.CalledProcessError as e:
        print_colored(f"错误: 创建venv环境失败: {e}", Colors.RED)
        return False

def list_conda_envs():
    """列出所有conda环境"""
    print_colored("所有Conda环境:", Colors.BLUE)
    print()
    
    if not check_command_exists("conda"):
        print_colored("错误: 未找到conda命令", Colors.RED)
        return False
    
    output, code = run_command(["conda", "env", "list"], check=False)
    if code == 0:
        print(output)
        print()
        current_env = os.environ.get("CONDA_DEFAULT_ENV")
        print_colored("当前激活的环境:", Colors.YELLOW)
        if current_env:
            print(f"  {current_env}")
        else:
            print("  无")
        return True
    else:
        print_colored(f"错误: {output}", Colors.RED)
        return False

def list_venv_envs():
    """列出venv环境"""
    print_colored("当前目录的Venv环境:", Colors.BLUE)
    print()
    
    current_dir = Path(".")
    venv_dirs = []
    
    for item in current_dir.iterdir():
        if item.is_dir():
            activate_script = None
            if (item / "bin" / "activate").exists():
                activate_script = item / "bin" / "activate"
            elif (item / "Scripts" / "activate.bat").exists():
                activate_script = item / "Scripts" / "activate.bat"
            
            if activate_script:
                venv_dirs.append(item.name)
    
    if venv_dirs:
        for venv_dir in venv_dirs:
            print(f"  - {venv_dir}")
        print()
        print_colored(f"找到 {len(venv_dirs)} 个venv环境", Colors.GREEN)
    else:
        print_colored("  未找到venv环境", Colors.YELLOW)
    
    return True

def show_conda_env(env_name: str):
    """显示conda环境详细信息"""
    if not env_name:
        print_colored("错误: 请提供环境名称", Colors.RED)
        return False
    
    if not check_command_exists("conda"):
        print_colored("错误: 未找到conda命令", Colors.RED)
        return False
    
    print_colored(f"Conda环境详细信息: {env_name}", Colors.BLUE)
    print()
    
    # 检查环境是否存在
    output, code = run_command(["conda", "env", "list"], check=False)
    if code != 0:
        print_colored(f"错误: {output}", Colors.RED)
        return False
    
    env_exists = False
    env_path = None
    for line in output.split("\n"):
        if line.strip().startswith(env_name):
            parts = line.split()
            if len(parts) >= 2:
                env_path = parts[-1]
                env_exists = True
                break
    
    if not env_exists:
        print_colored(f"错误: 环境 '{env_name}' 不存在", Colors.RED)
        return False
    
    print_colored("环境路径:", Colors.GREEN)
    print(f"  {env_path}")
    print()
    
    # 列出已安装的包
    print_colored("已安装的包:", Colors.GREEN)
    output, code = run_command(["conda", "list", "-n", env_name], check=False)
    if code == 0:
        lines = output.split("\n")[:22]  # 前20个包 + 2行标题
        print("\n".join(lines))
        print()
        print_colored("（显示前20个包，使用 'conda list -n {env_name}' 查看全部）", Colors.YELLOW)
    else:
        print_colored(f"  无法获取包列表: {output}", Colors.RED)
    
    return True

def show_venv_env(env_path: str):
    """显示venv环境详细信息"""
    if not env_path:
        print_colored("错误: 请提供环境路径", Colors.RED)
        return False
    
    env_dir = Path(env_path)
    if not env_dir.exists():
        print_colored(f"错误: 路径 '{env_path}' 不存在", Colors.RED)
        return False
    
    # 检查是否是有效的venv环境
    python_cmd = None
    if (env_dir / "bin" / "python").exists():
        python_cmd = env_dir / "bin" / "python"
    elif (env_dir / "Scripts" / "python.exe").exists():
        python_cmd = env_dir / "Scripts" / "python.exe"
    else:
        print_colored(f"错误: '{env_path}' 不是有效的venv环境", Colors.RED)
        return False
    
    print_colored(f"Venv环境详细信息: {env_path}", Colors.BLUE)
    print()
    
    print_colored("环境路径:", Colors.GREEN)
    print(f"  {env_dir.resolve()}")
    print()
    
    # 获取Python版本
    print_colored("Python版本:", Colors.GREEN)
    output, code = run_command([str(python_cmd), "--version"], check=False)
    if code == 0:
        print(f"  {output}")
    else:
        print("  无法获取版本")
    print()
    
    # 列出已安装的包
    print_colored("已安装的包:", Colors.GREEN)
    output, code = run_command([str(python_cmd), "-m", "pip", "list"], check=False)
    if code == 0:
        lines = output.split("\n")[:22]  # 前20个包 + 2行标题
        print("\n".join(lines))
        print()
        print_colored(f"（显示前20个包，使用 '{python_cmd} -m pip list' 查看全部）", Colors.YELLOW)
    else:
        print_colored(f"  无法获取包列表: {output}", Colors.RED)
    
    return True

def install_deps(env_type: str, env_name: str):
    """安装项目依赖"""
    if not env_type or not env_name:
        print_colored("错误: 请提供环境类型和环境名称", Colors.RED)
        print("用法: python manage_env.py install-deps <conda|venv> <env_name>")
        return False
    
    print_colored(f"安装项目依赖到 {env_type} 环境: {env_name}", Colors.GREEN)
    print()
    
    # 检查依赖文件
    if not Path("requirements.txt").exists():
        print_colored("警告: 未找到 requirements.txt", Colors.YELLOW)
    
    if not Path("pyproject.toml").exists():
        print_colored("警告: 未找到 pyproject.toml", Colors.YELLOW)
    
    if env_type == "conda":
        if not check_command_exists("conda"):
            print_colored("错误: 未找到conda命令", Colors.RED)
            return False
        
        print_colored("激活conda环境并安装依赖...", Colors.BLUE)
        print()
        print("执行以下命令:")
        print(f"  conda activate {env_name}")
        print('  pip install -e ".[torch,metrics]" --no-build-isolation')
        if Path("requirements.txt").exists():
            print("  pip install -r requirements.txt")
        return True
    
    elif env_type == "venv":
        env_dir = Path(env_name)
        if not env_dir.exists():
            print_colored(f"错误: venv环境 '{env_name}' 不存在", Colors.RED)
            return False
        
        # 确定pip命令路径
        pip_cmd = None
        if (env_dir / "bin" / "pip").exists():
            pip_cmd = env_dir / "bin" / "pip"
        elif (env_dir / "Scripts" / "pip.exe").exists():
            pip_cmd = env_dir / "Scripts" / "pip.exe"
        else:
            print_colored("错误: 未找到pip命令", Colors.RED)
            return False
        
        print_colored("使用venv环境的pip安装依赖...", Colors.BLUE)
        
        # 升级pip
        run_command([str(pip_cmd), "install", "--upgrade", "pip"], check=False)
        
        # 安装项目依赖
        if Path("pyproject.toml").exists():
            output, code = run_command([
                str(pip_cmd), "install", "-e", ".[torch,metrics]", 
                "--no-build-isolation"
            ], check=False)
            if code != 0:
                print_colored(f"警告: 安装项目依赖失败: {output}", Colors.YELLOW)
        
        if Path("requirements.txt").exists():
            output, code = run_command([
                str(pip_cmd), "install", "-r", "requirements.txt"
            ], check=False)
            if code != 0:
                print_colored(f"警告: 安装requirements.txt失败: {output}", Colors.YELLOW)
        
        print_colored("✓ 依赖安装完成", Colors.GREEN)
        return True
    
    else:
        print_colored(f"错误: 不支持的环境类型 '{env_type}'", Colors.RED)
        print("支持的类型: conda, venv")
        return False

def show_help():
    """显示帮助信息"""
    print_colored("环境管理脚本", Colors.BLUE)
    print()
    print("用法: python manage_env.py [命令] [选项]")
    print()
    print("命令:")
    print("  create-conda <env_name> [python_version]  创建conda环境")
    print("  create-venv <env_name>                    创建venv虚拟环境")
    print("  list-conda                                列出所有conda环境")
    print("  list-venv                                 列出当前目录的venv环境")
    print("  show-conda <env_name>                     显示conda环境详细信息")
    print("  show-venv <env_path>                      显示venv环境详细信息")
    print("  install-deps <env_type> <env_name>        安装项目依赖")
    print("  help                                       显示此帮助信息")
    print()
    print("示例:")
    print("  python manage_env.py create-conda llamafactory 3.10")
    print("  python manage_env.py list-conda")
    print("  python manage_env.py show-conda llamafactory")
    print("  python manage_env.py install-deps conda llamafactory")

def main():
    """主函数"""
    if len(sys.argv) < 2:
        show_help()
        return
    
    command = sys.argv[1]
    
    if command == "create-conda":
        env_name = sys.argv[2] if len(sys.argv) > 2 else ""
        python_version = sys.argv[3] if len(sys.argv) > 3 else "3.10"
        create_conda_env(env_name, python_version)
    
    elif command == "create-venv":
        env_name = sys.argv[2] if len(sys.argv) > 2 else ""
        create_venv_env(env_name)
    
    elif command == "list-conda":
        list_conda_envs()
    
    elif command == "list-venv":
        list_venv_envs()
    
    elif command == "show-conda":
        env_name = sys.argv[2] if len(sys.argv) > 2 else ""
        show_conda_env(env_name)
    
    elif command == "show-venv":
        env_path = sys.argv[2] if len(sys.argv) > 2 else ""
        show_venv_env(env_path)
    
    elif command == "install-deps":
        env_type = sys.argv[2] if len(sys.argv) > 2 else ""
        env_name = sys.argv[3] if len(sys.argv) > 3 else ""
        install_deps(env_type, env_name)
    
    elif command in ["help", "--help", "-h"]:
        show_help()
    
    else:
        print_colored(f"错误: 未知命令 '{command}'", Colors.RED)
        print()
        show_help()
        sys.exit(1)

if __name__ == "__main__":
    main()