#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Python 环境列表更新脚本
功能：对比虚拟环境中已安装的包和 environment.txt，如果不一致则更新 environment.txt
"""

import os
import sys
import subprocess
import platform
from pathlib import Path

# 获取脚本所在目录
SCRIPT_DIR = Path(__file__).parent.absolute()
VENV_PATH = SCRIPT_DIR / "env"
ENVIRONMENT_FILE = SCRIPT_DIR / "environment.txt"

# 根据操作系统确定虚拟环境的 pip 路径
if platform.system() == "Windows":
    VENV_PIP = VENV_PATH / "Scripts" / "pip.exe"
else:
    VENV_PIP = VENV_PATH / "bin" / "pip"


def run_command(cmd, check=True, capture_output=True):
    """运行命令并返回结果"""
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            check=check,
            capture_output=capture_output,
            text=True,
            encoding='utf-8'
        )
        return result.returncode == 0, result.stdout, result.stderr
    except subprocess.CalledProcessError as e:
        return False, e.stdout if hasattr(e, 'stdout') else "", str(e)


def get_installed_packages():
    """获取虚拟环境中已安装的包列表（使用 pip freeze）"""
    if not VENV_PATH.exists():
        print("[ERROR] Virtual environment not found")
        return None
    
    cmd = f'"{VENV_PIP}" freeze'
    success, output, error = run_command(cmd, check=False)
    
    if not success:
        print(f"[ERROR] Failed to get installed packages: {error}")
        return None
    
    # 解析输出，提取包名和版本
    packages = {}
    for line in output.strip().split('\n'):
        line = line.strip()
        if line and not line.startswith('#'):
            # 格式：package==version 或 package>=version 等
            if '==' in line:
                parts = line.split('==', 1)
                packages[parts[0].strip().lower()] = line.strip()
            else:
                # 如果没有版本号，使用整行
                pkg_name = line.split('>=')[0].split('<=')[0].split('>')[0].split('<')[0].split('~=')[0].strip()
                packages[pkg_name.lower()] = line.strip()
    
    return packages


def read_environment_file():
    """读取 environment.txt 中的包列表"""
    if not ENVIRONMENT_FILE.exists():
        print(f"[WARN] {ENVIRONMENT_FILE} not found, will create new one")
        return {}
    
    packages = {}
    with open(ENVIRONMENT_FILE, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#'):
                # 提取包名（支持各种版本操作符）
                if '==' in line:
                    parts = line.split('==', 1)
                    packages[parts[0].strip().lower()] = line.strip()
                else:
                    pkg_name = line.split('>=')[0].split('<=')[0].split('>')[0].split('<')[0].split('~=')[0].strip()
                    packages[pkg_name.lower()] = line.strip()
    
    return packages


def compare_and_update():
    """对比并更新 environment.txt"""
    print("Comparing installed packages with environment.txt...")
    print("=" * 60)
    
    # 获取已安装的包
    installed_packages = get_installed_packages()
    if installed_packages is None:
        sys.exit(1)
    
    # 读取 environment.txt 中的包
    file_packages = read_environment_file()
    
    # 对比差异
    installed_set = set(installed_packages.keys())
    file_set = set(file_packages.keys())
    
    added_packages = installed_set - file_set
    removed_packages = file_set - installed_set
    changed_packages = []
    
    # 检查版本变化
    for pkg_name in installed_set & file_set:
        if installed_packages[pkg_name] != file_packages[pkg_name]:
            changed_packages.append(pkg_name)
    
    # 显示差异
    if added_packages:
        print(f"\n[+] Added packages ({len(added_packages)}):")
        for pkg in sorted(added_packages):
            print(f"  + {installed_packages[pkg]}")
    
    if removed_packages:
        print(f"\n[-] Removed packages ({len(removed_packages)}):")
        for pkg in sorted(removed_packages):
            print(f"  - {file_packages[pkg]}")
    
    if changed_packages:
        print(f"\n[~] Changed packages ({len(changed_packages)}):")
        for pkg in sorted(changed_packages):
            print(f"  ~ {file_packages[pkg]} -> {installed_packages[pkg]}")
    
    # 判断是否需要更新
    if not added_packages and not removed_packages and not changed_packages:
        print("\n[OK] environment.txt is up to date")
        print(f"     Total packages: {len(installed_packages)}")
        return True
    
    # 更新 environment.txt
    print(f"\nUpdating {ENVIRONMENT_FILE}...")
    
    # 使用 pip freeze 获取完整列表（包含所有依赖）
    cmd = f'"{VENV_PIP}" freeze'
    success, output, error = run_command(cmd, check=False)
    
    if not success:
        print(f"[ERROR] Failed to get installed packages: {error}")
        sys.exit(1)
    
    # 写入文件（使用 UTF-8 无 BOM 编码）
    with open(ENVIRONMENT_FILE, 'w', encoding='utf-8', newline='\n') as f:
        f.write(output)
    
    # 统计包数量
    package_count = len([line for line in output.strip().split('\n') if line.strip()])
    
    print(f"[OK] {ENVIRONMENT_FILE} updated successfully")
    print(f"     Total packages: {package_count}")
    
    return True


def main():
    """主函数"""
    if not compare_and_update():
        sys.exit(1)
    sys.exit(0)


if __name__ == "__main__":
    main()