#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Python 依赖安装和同步脚本
功能：检查、安装 Python 依赖到虚拟环境，然后同步所有已安装的包到 environment.txt
约定：由嵌入式 Python（python/x64/py 或 python/arm64/py）运行；虚拟环境在 env，可用 virtualenv 创建。
"""

import os
import sys
import subprocess
import platform
import tempfile
import urllib.request
import urllib.error
import ssl
from pathlib import Path

# pip 损坏时用 get-pip.py 重装；优先国内源避免 SSL/外网问题（见 https://pip.pypa.io/en/stable/installation/）
GET_PIP_URLS = [
    "http://mirrors.aliyun.com/pypi/get-pip.py",
    "https://mirrors.aliyun.com/pypi/get-pip.py",
    "https://bootstrap.pypa.io/get-pip.py",
]

# 国内 pip 镜像源，避免外网慢或超时（可改为阿里/豆瓣等）
PIP_INDEX_URL = os.environ.get("PIP_INDEX_URL", "https://pypi.tuna.tsinghua.edu.cn/simple")
PIP_TRUSTED_HOST = "pypi.tuna.tsinghua.edu.cn"
PIP_INDEX_ARGS = f'-i {PIP_INDEX_URL} --trusted-host {PIP_TRUSTED_HOST}'

# 脚本所在目录即本环境目录（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 路径
# Windows 优先使用 .bat 脚本（可读、可编辑），否则用 .exe
if platform.system() == "Windows":
    VENV_PYTHON = VENV_PATH / "Scripts" / "python.exe"
    VENV_PIP = VENV_PATH / "Scripts" / "pip.exe"
    PY_SCRIPTS = SCRIPT_DIR / "py" / "Scripts"
    PY_PIP = PY_SCRIPTS / "pip.bat" if (PY_SCRIPTS / "pip.bat").exists() else PY_SCRIPTS / "pip.exe"
    PY_VIRTUALENV = PY_SCRIPTS / "virtualenv.bat" if (PY_SCRIPTS / "virtualenv.bat").exists() else PY_SCRIPTS / "virtualenv.exe"
else:
    VENV_PYTHON = VENV_PATH / "bin" / "python"
    VENV_PIP = VENV_PATH / "bin" / "pip"
    PY_SCRIPTS = SCRIPT_DIR / "py" / "bin"
    PY_PIP = PY_SCRIPTS / "pip"
    PY_VIRTUALENV = PY_SCRIPTS / "virtualenv"


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 _find_or_download_get_pip(get_pip_path):
    """优先使用项目内已有 get-pip.py，否则从国内源/官方源下载"""
    for local_dir in (SCRIPT_DIR, PROJECT_ROOT):
        candidate = Path(local_dir) / "get-pip.py"
        if candidate.exists():
            with open(candidate, "rb") as f:
                Path(get_pip_path).write_bytes(f.read())
            return True
    return _download_get_pip(get_pip_path)


def _download_get_pip(get_pip_path):
    """从国内源或官方源下载 get-pip.py，HTTPS 遇 SSL 时用不验证上下文重试"""
    last_err = None
    for url in GET_PIP_URLS:
        try:
            ctx = None
            if url.startswith('https://'):
                ctx = ssl.create_default_context()
            req = urllib.request.urlopen(url, timeout=30, context=ctx)
            with open(get_pip_path, 'wb') as f:
                f.write(req.read())
            return True
        except urllib.error.URLError as e:
            last_err = e
            if 'SSL' in str(e) or 'CERTIFICATE' in str(e).upper():
                try:
                    ctx = ssl._create_unverified_context()
                    req = urllib.request.urlopen(url, timeout=30, context=ctx)
                    with open(get_pip_path, 'wb') as f:
                        f.write(req.read())
                    return True
                except Exception as retry_e:
                    last_err = retry_e
        except Exception as e:
            last_err = e
    if last_err:
        print(f"[WARN] Failed to download get-pip.py: {last_err}")
    return False


def repair_pip_with_get_pip():
    """当 pip 损坏（如缺少 pip._internal.operations.build）时，用 get-pip.py 强制重装 pip（优先国内源）"""
    try:
        fd, get_pip_path = tempfile.mkstemp(suffix='.py')
        os.close(fd)
        if not _find_or_download_get_pip(get_pip_path):
            return False
        ok, out, err = run_command(
            f'"{sys.executable}" "{get_pip_path}" --force-reinstall -i {PIP_INDEX_URL} --trusted-host {PIP_TRUSTED_HOST} --no-warn-script-location',
            check=False
        )
        try:
            os.unlink(get_pip_path)
        except Exception:
            pass
        return ok
    except Exception as e:
        print(f"[WARN] Failed to download or run get-pip.py: {e}")
        return False


def check_pip():
    """先检测 pip 是否已安装，未安装则退出"""
    if not PY_PIP.exists():
        print(f"[X] pip is missing. Put pip.exe in: {PY_SCRIPTS}")
        sys.exit(1)
    print("[OK] pip is ready")


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, out, err = run_command(f'"{sys.executable}" -m venv "{VENV_PATH}"', check=False)
        merged_err = (err or "") + (out or "")
        if not success and "No module named venv" in merged_err:
            # 嵌入式 Python 无 venv，优先使用已有的 virtualenv.exe，否则用 pip 安装
            if PY_VIRTUALENV.exists():
                venv_cmd = f'"{PY_VIRTUALENV}" "{VENV_PATH}"'
            else:
                if not PY_PIP.exists():
                    print(f"[X] Missing: {PY_PIP} and {PY_VIRTUALENV}. Put pip and virtualenv in py/Scripts/")
                    sys.exit(1)
                print("[WARN] venv module is missing. Installing virtualenv with pip...")
                pip_ok, pip_out, pip_err = run_command(f'"{PY_PIP}" install {PIP_INDEX_ARGS} --no-warn-script-location virtualenv', check=False)
                if not pip_ok:
                    print(f"[X] Failed to install virtualenv. Error: {(pip_err or '') + (pip_out or '')}")
                    sys.exit(1)
                print("[OK] virtualenv installed successfully via pip")
                venv_cmd = f'"{PY_VIRTUALENV}" "{VENV_PATH}"' if PY_VIRTUALENV.exists() else f'"{sys.executable}" -m virtualenv "{VENV_PATH}"'
            success, out, err = run_command(venv_cmd, check=False)
            merged_err = (err or "") + (out or "")

        # virtualenv 损坏（如 No discovery plugin found）时：先删原虚拟环境，再重装 virtualenv 并重建
        if not success:
            err_msg = merged_err.strip() or '(no error output)'
            if "No discovery plugin found" in err_msg or "discovery plugin" in err_msg.lower():
                import shutil
                # 先删除原来的虚拟环境（含残缺的 env 目录）
                if VENV_PATH.exists():
                    try:
                        shutil.rmtree(VENV_PATH)
                        print("[OK] Old virtual environment removed, recreating...")
                    except Exception as e:
                        print(f"[WARN] Failed to remove env: {e}")
                print("[WARN] virtualenv is broken, reinstalling virtualenv...")
                reinstall_ok, _, reinstall_err = run_command(
                    f'"{sys.executable}" -m pip install {PIP_INDEX_ARGS} --no-warn-script-location --force-reinstall virtualenv',
                    check=False
                )
                # pip 自身损坏（如缺少 pip._internal.operations.build）时用 get-pip.py 重装（不依赖残缺 pip，优先国内源）
                if not reinstall_ok and ("pip._internal" in (reinstall_err or "") or "ModuleNotFoundError" in (reinstall_err or "")):
                    print("[WARN] pip is broken (embedded pip may be incomplete), reinstalling pip with get-pip.py...")
                    if repair_pip_with_get_pip():
                        reinstall_ok, _, reinstall_err = run_command(
                            f'"{sys.executable}" -m pip install {PIP_INDEX_ARGS} --no-warn-script-location --force-reinstall virtualenv',
                            check=False
                        )
                    if not reinstall_ok:
                        ensure_ok, _, _ = run_command(f'"{sys.executable}" -m ensurepip --upgrade', check=False)
                        if ensure_ok:
                            reinstall_ok, _, reinstall_err = run_command(
                                f'"{sys.executable}" -m pip install {PIP_INDEX_ARGS} --no-warn-script-location --force-reinstall virtualenv',
                                check=False
                            )
                    if not reinstall_ok:
                        print(f"[X] Failed to reinstall virtualenv: {reinstall_err or '(no output)'}")
                        print("[HINT] Embedded pip in python/x64/py may be incomplete (missing pip._internal.operations.build).")
                        print("  Option 1: put get-pip.py in python/x64/ or project root, then rerun this script (local file is preferred).")
                        print("  Option 2: run manually (download get-pip.py from: http://mirrors.aliyun.com/pypi/get-pip.py):")
                        print("  python get-pip.py --force-reinstall -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn")
                        sys.exit(1)
                if not reinstall_ok:
                    print(f"[X] Failed to reinstall virtualenv: {reinstall_err or '(no output)'}")
                    sys.exit(1)
                print("[OK] virtualenv reinstalled")
                # 重新创建虚拟环境：优先 -m venv，再 fallback 到 virtualenv
                success, out, err = run_command(f'"{sys.executable}" -m venv "{VENV_PATH}"', check=False)
                merged_err = (err or "") + (out or "")
                if not success and "No module named venv" in merged_err:
                    venv_cmd = f'"{PY_VIRTUALENV}" "{VENV_PATH}"' if PY_VIRTUALENV.exists() else f'"{sys.executable}" -m virtualenv "{VENV_PATH}"'
                    success, out, err = run_command(venv_cmd, check=False)
                    merged_err = (err or "") + (out or "")

        if not success:
            err_msg = merged_err.strip() or '(no error output)'
            print(f"[X] Failed to create virtual environment: {err_msg}")
            if "No discovery plugin found" in err_msg:
                print("[HINT] virtualenv dist-info is missing entry_points.txt. Check virtualenv-*.dist-info/")
            if "pythonw.exe" in err_msg or "FileNotFoundError" in err_msg:
                print("[HINT] py/pythonw.exe is required. You can copy it from python.exe.")
            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 = []
    total = len(packages)
    for i, package in enumerate(packages, 1):
        pkg_display = package.split("==")[0].strip() if "==" in package else package.strip()
        print(f"[{i}/{total}] Installing {pkg_display} ...", end=" ", flush=True)
        cmd = f'"{venv_pip}" install {PIP_INDEX_ARGS} --no-warn-script-location {package}'
        success, _, error = run_command(cmd, check=False)
        if success:
            print("OK")
        else:
            print("FAILED")
            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():
    """主函数"""
    # 1. 先检测 pip 是否安装
    check_pip()
    # 2. 再检测并创建虚拟环境
    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(f"\nInstalling {missing_count} missing package(s)...\n")
        
        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()