#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Python 依赖安装和同步脚本
功能：检查、安装 Python 依赖到虚拟环境，然后同步所有已安装的包到 environment.txt
"""

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

# 脚本所在目录即本环境目录（python/arm64 或 python/x64），venv 固定为 env
SCRIPT_DIR = Path(__file__).parent.absolute()
PROJECT_ROOT = SCRIPT_DIR.parent.parent.absolute()
_env_path = os.environ.get("PYTHON_VENV_PATH", "").strip()
VENV_PATH = Path(_env_path) if _env_path else SCRIPT_DIR / "env"
ENVIRONMENT_FILE = SCRIPT_DIR / "environment.txt"
REQUIREMENTS_FILE = PROJECT_ROOT / "requirements.txt"

# 根据操作系统确定虚拟环境的 Python 和 pip 路径
if platform.system() == "Windows":
    VENV_PYTHON = VENV_PATH / "Scripts" / "python.exe"
    VENV_PIP = VENV_PATH / "Scripts" / "pip.exe"
else:
    VENV_PYTHON = VENV_PATH / "bin" / "python"
    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 _venv_home():
    """读取 venv 的 pyvenv.cfg 中的 home（创建该 venv 的 Python 路径）"""
    cfg = VENV_PATH / "pyvenv.cfg"
    if not cfg.exists():
        return None
    try:
        for line in cfg.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line.startswith("home ") or line.startswith("home="):
                return line.split("=", 1)[-1].strip()
    except Exception:
        pass
    return None


def ensure_venv():
    """确保虚拟环境存在，且由当前 Python (sys.executable) 创建"""
    current_python_home = str(Path(sys.executable).resolve().parent)
    if VENV_PATH.exists():
        existing_home = _venv_home()
        if existing_home:
            existing_home = str(Path(existing_home).resolve())
        if existing_home and existing_home != current_python_home:
            print("[WARN] venv was created by another Python, recreating with current Python...")
            import shutil
            try:
                shutil.rmtree(VENV_PATH)
            except Exception as e:
                print(f"[X] Failed to remove old venv: {e}")
                sys.exit(1)
        elif existing_home == current_python_home:
            return True
    if not VENV_PATH.exists():
        print("[WARN] Virtual environment not found, creating...")
        success, _, error = run_command(f'"{sys.executable}" -m venv "{VENV_PATH}"', check=False)
        if not success and "No module named venv" in (error or ""):
            print("[WARN] venv module not found, using virtualenv...")
            run_command(f'"{sys.executable}" -m pip install virtualenv -q', check=False)
            success, _, error = run_command(f'"{sys.executable}" -m virtualenv "{VENV_PATH}"', check=False)
        if not success:
            print(f"[X] Failed to create virtual environment: {error}")
            sys.exit(1)
        print("[OK] Virtual environment created successfully")
    return True


def get_venv_pip():
    """获取虚拟环境的 pip 命令"""
    if platform.system() == "Windows":
        return str(VENV_PIP)
    else:
        return str(VENV_PIP)


def read_dependencies(source_file):
    """读取依赖列表"""
    if not source_file.exists():
        return []
    
    dependencies = []
    with open(source_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            # 跳过注释和空行
            if line and not line.startswith('#'):
                dependencies.append(line)
    return dependencies


def get_installed_packages_from_filesystem():
    """直接从文件系统获取已安装的包列表（快速方法）"""
    if platform.system() == "Windows":
        site_packages = VENV_PATH / "Lib" / "site-packages"
    else:
        # Linux/Mac: 需要找到 site-packages 路径
        import sysconfig
        site_packages = Path(sysconfig.get_path('purelib', vars={'base': str(VENV_PATH)}))
    
    installed_packages = set()
    
    if site_packages.exists():
        for item in site_packages.iterdir():
            if item.is_dir():
                pkg_name = item.name
                
                # 处理 .dist-info 和 .egg-info 文件夹（最准确的包名来源）
                if pkg_name.endswith('.dist-info'):
                    # 从 dist-info 文件夹名提取包名（格式：package-name-version.dist-info）
                    parts = pkg_name.replace('.dist-info', '').rsplit('-', 1)
                    if len(parts) >= 1:
                        installed_packages.add(parts[0].lower().replace('_', '-'))
                elif pkg_name.endswith('.egg-info'):
                    # 从 egg-info 文件夹名提取包名
                    parts = pkg_name.replace('.egg-info', '').rsplit('-', 1)
                    if len(parts) >= 1:
                        installed_packages.add(parts[0].lower().replace('_', '-'))
                elif pkg_name not in ['__pycache__', 'dist-info', 'egg-info']:
                    # 检查是否是 Python 包（有 __init__.py 或 .py 文件）
                    if (item / "__init__.py").exists() or any(item.glob("*.py")):
                        pkg_lower = pkg_name.lower()
                        installed_packages.add(pkg_lower)
                        # 添加下划线和连字符的变体
                        installed_packages.add(pkg_lower.replace('_', '-'))
                        installed_packages.add(pkg_lower.replace('-', '_'))
        
        # 特殊映射：opencv-python 安装后显示为 cv2
        if 'cv2' in installed_packages:
            installed_packages.add('opencv-python')
            installed_packages.add('opencv-contrib-python')
            installed_packages.add('opencv-python-headless')
    
    return installed_packages


def check_package_installed(package_name, installed_packages_set=None):
    """检查包是否已安装（使用文件系统快速检查）"""
    # 提取包名（支持 ==, >=, <=, >, <, ~= 等版本操作符）
    pkg_name = package_name.split('==')[0].split('>=')[0].split('<=')[0].split('>')[0].split('<')[0].split('~=')[0].strip()
    pkg_name_lower = pkg_name.lower()
    
    # 如果没有提供已安装包集合，则获取一次（避免重复调用）
    if installed_packages_set is None:
        installed_packages_set = get_installed_packages_from_filesystem()
    
    # 快速检查（使用已获取的集合）
    return (
        pkg_name_lower in installed_packages_set or
        pkg_name_lower.replace('-', '_') in installed_packages_set or
        pkg_name_lower.replace('_', '-') in installed_packages_set
    )


def install_packages(packages, source_file, venv_pip):
    """安装包到虚拟环境"""
    failed_packages = []
    
    if source_file == REQUIREMENTS_FILE and REQUIREMENTS_FILE.exists():
        # 使用 requirements.txt 批量安装
        cmd = f'"{venv_pip}" install -r "{source_file}"'
        success, _, error = run_command(cmd, check=False)
        if not success:
            print(f"[X] Installation failed: {error}")
            return False, failed_packages
    else:
        # 逐个安装
        for package in packages:
            cmd = f'"{venv_pip}" install {package}'
            success, _, error = run_command(cmd, check=False)
            if not success:
                print(f"[X] Failed to install: {package}")
                failed_packages.append(package)
    
    if failed_packages:
        return False, failed_packages
    return True, []


def sync_environment_file(venv_pip, silent=False):
    """同步所有已安装的包到 environment.txt"""
    cmd = f'"{venv_pip}" freeze'
    success, output, error = run_command(cmd, check=False)
    
    if not success:
        if not silent:
            print(f"[X] Failed to get installed packages list: {error}")
        return False
    
    # 使用 UTF-8 无 BOM 编码写入文件
    with open(ENVIRONMENT_FILE, 'w', encoding='utf-8', newline='\n') as f:
        f.write(output)
    
    if not silent:
        package_count = len([line for line in output.strip().split('\n') if line.strip()])
        print(f"[OK] All installed packages synced to {ENVIRONMENT_FILE}")
        print(f"      Total packages: {package_count}")
    return True


def main():
    """主函数"""
    # 确保虚拟环境存在
    if not ensure_venv():
        sys.exit(1)
    
    venv_pip = get_venv_pip()
    
    # 确定依赖源文件（优先使用 requirements.txt，如果没有则使用 environment.txt）
    if REQUIREMENTS_FILE.exists():
        source_file = REQUIREMENTS_FILE
    elif ENVIRONMENT_FILE.exists():
        source_file = ENVIRONMENT_FILE
    else:
        sync_environment_file(venv_pip)
        sys.exit(0)
    
    # 读取依赖列表
    required_packages = read_dependencies(source_file)
    
    if not required_packages:
        print("[OK] No dependencies specified")
        sys.exit(0)
    
    # 快速检查缺失的依赖（使用文件系统）
    missing_packages = []
    installed_count = 0
    missing_count = 0
    
    # 一次性获取所有已安装的包（只检查一次文件系统）
    installed_packages_set = get_installed_packages_from_filesystem()
    
    for package in required_packages:
        package_line = package.strip()
        if not package_line:
            continue
        
        # 提取包名
        package_name = package_line.split('==')[0].split('>=')[0].split('<=')[0].split('>')[0].split('<')[0].split('~=')[0].strip()
        pkg_name_lower = package_name.lower()
        
        # 快速检查（使用已获取的集合）
        is_installed = (
            pkg_name_lower in installed_packages_set or
            pkg_name_lower.replace('-', '_') in installed_packages_set or
            pkg_name_lower.replace('_', '-') in installed_packages_set
        )
        
        if is_installed:
            installed_count += 1
        else:
            missing_packages.append(package_line)
            missing_count += 1
    
    # 如果有缺失的依赖，显示必要信息并安装
    if missing_count > 0:
        print(f"[X] Missing {missing_count} package(s) out of {len(required_packages)}")
        
        print("Missing packages:")
        for missing in missing_packages:
            print(f"  - {missing}")
        
        print("\nInstalling missing packages...")
        
        success, failed = install_packages(missing_packages, source_file, venv_pip)
        if success:
            print("[OK] All packages installed successfully")
        else:
            if failed:
                print(f"[X] Failed to install {len(failed)} package(s):")
                for pkg in failed:
                    print(f"  - {pkg}")
            else:
                print("[X] Some packages installation failed")
            # 即使有失败，也继续同步已安装的包
            print("[WARN] Continuing to sync installed packages...")
        
        # 同步所有已安装的包到 environment.txt
        sync_environment_file(venv_pip)
    else:
        # 所有依赖都齐全时，只显示一行信息（包含同步结果）
        sync_environment_file(venv_pip, silent=True)
        print(f"[OK] All dependencies are installed ({len(required_packages)} packages)")
    
    sys.exit(0)


if __name__ == "__main__":
    main()