"""
环境检查模块
用于检查并自动安装 UniLabOS 运行所需的 Python 包
"""

import argparse
import importlib
import locale
import shutil
import subprocess
import sys
from pathlib import Path
from typing import List, Optional

from unilabos.utils.banner_print import print_status


# ---------------------------------------------------------------------------
# 底层安装工具
# ---------------------------------------------------------------------------

def _is_chinese_locale() -> bool:
    try:
        lang = locale.getdefaultlocale()[0]
        return bool(lang and ("zh" in lang.lower() or "chinese" in lang.lower()))
    except Exception:
        return False


_USE_UV: Optional[bool] = None


def _has_uv() -> bool:
    global _USE_UV
    if _USE_UV is None:
        _USE_UV = shutil.which("uv") is not None
    return _USE_UV


def _install_packages(
    packages: List[str],
    upgrade: bool = False,
    label: str = "",
) -> bool:
    """
    安装/升级一组包。优先 uv pip install，回退 sys pip。
    逐个安装，任意一个失败不影响后续包。

    Returns:
        True if all succeeded, False otherwise.
    """
    if not packages:
        return True

    is_chinese = _is_chinese_locale()
    use_uv = _has_uv()
    failed: List[str] = []

    for pkg in packages:
        action_word = "升级" if upgrade else "安装"
        if label:
            print_status(f"[{label}] 正在{action_word} {pkg}...", "info")
        else:
            print_status(f"正在{action_word} {pkg}...", "info")

        if use_uv:
            cmd = ["uv", "pip", "install"]
            if upgrade:
                cmd.append("--upgrade")
            cmd.append(pkg)
            if is_chinese:
                cmd.extend(["--index-url", "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"])
        else:
            cmd = [sys.executable, "-m", "pip", "install"]
            if upgrade:
                cmd.append("--upgrade")
            cmd.append(pkg)
            if is_chinese:
                cmd.extend(["-i", "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"])

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                installer = "uv" if use_uv else "pip"
                print_status(f"✓ {pkg} {action_word}成功 (via {installer})", "success")
            else:
                stderr_short = result.stderr.strip().split("\n")[-1] if result.stderr else "unknown error"
                print_status(f"× {pkg} {action_word}失败: {stderr_short}", "error")
                failed.append(pkg)
        except subprocess.TimeoutExpired:
            print_status(f"× {pkg} {action_word}超时 (300s)", "error")
            failed.append(pkg)
        except Exception as e:
            print_status(f"× {pkg} {action_word}异常: {e}", "error")
            failed.append(pkg)

    if failed:
        print_status(f"有 {len(failed)} 个包操作失败: {', '.join(failed)}", "error")
        return False
    return True


# ---------------------------------------------------------------------------
# requirements.txt 安装（可多次调用）
# ---------------------------------------------------------------------------

def install_requirements_txt(req_path: str | Path, label: str = "") -> bool:
    """
    读取一个 requirements.txt 文件，检查缺失的包并安装。

    Args:
        req_path: requirements.txt 文件路径
        label: 日志前缀标签（如 "device_package_sim"）

    Returns:
        True if all ok, False if any install failed.
    """
    req_path = Path(req_path)
    if not req_path.exists():
        return True

    tag = label or req_path.parent.name
    print_status(f"[{tag}] 检查依赖: {req_path}", "info")

    reqs: List[str] = []
    with open(req_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and not line.startswith("-"):
                reqs.append(line)

    if not reqs:
        return True

    missing: List[str] = []
    for req in reqs:
        pkg_import = req.split(">=")[0].split("==")[0].split("<")[0].split("[")[0].split(">")[0].strip()
        pkg_import = pkg_import.replace("-", "_")
        try:
            importlib.import_module(pkg_import)
        except ImportError:
            missing.append(req)

    if not missing:
        print_status(f"[{tag}] ✓ 依赖检查通过 ({len(reqs)} 个包)", "success")
        return True

    print_status(f"[{tag}] 缺失 {len(missing)} 个依赖: {', '.join(missing)}", "warning")
    return _install_packages(missing, label=tag)


def check_device_package_requirements(devices_dirs: list[str]) -> bool:
    """
    检查 --devices 指定的所有外部设备包目录中的 requirements.txt。
    对每个目录查找 requirements.txt（先在目录内找，再在父目录找）。
    """
    if not devices_dirs:
        return True

    all_ok = True
    for d in devices_dirs:
        d_path = Path(d).resolve()
        req_file = d_path / "requirements.txt"
        if not req_file.exists():
            req_file = d_path.parent / "requirements.txt"
        if not req_file.exists():
            continue
        if not install_requirements_txt(req_file, label=d_path.name):
            all_ok = False

    return all_ok


# ---------------------------------------------------------------------------
# UniLabOS 核心环境检查
# ---------------------------------------------------------------------------

class EnvironmentChecker:
    """环境检查器"""

    def __init__(self):
        self.required_packages = {
            "websockets": "websockets",
            "msgcenterpy": "msgcenterpy",
            "orjson": "orjson",
            "opentrons_shared_data": "opentrons_shared_data",
            "typing_extensions": "typing_extensions",
            "crcmod": "crcmod-plus",
        }

        self.special_packages = {"pylabrobot": "git+https://github.com/Xuwznln/pylabrobot.git"}

        self.version_requirements = {
            "msgcenterpy": "0.1.8",
        }

        self.missing_packages: List[tuple] = []
        self.failed_installs: List[tuple] = []
        self.packages_need_upgrade: List[tuple] = []

    def check_package_installed(self, package_name: str) -> bool:
        try:
            importlib.import_module(package_name)
            return True
        except ImportError:
            return False

    def get_package_version(self, package_name: str) -> str | None:
        try:
            module = importlib.import_module(package_name)
            return getattr(module, "__version__", None)
        except (ImportError, AttributeError):
            return None

    def compare_version(self, current: str, required: str) -> bool:
        try:
            current_parts = [int(x) for x in current.split(".")]
            required_parts = [int(x) for x in required.split(".")]
            max_len = max(len(current_parts), len(required_parts))
            current_parts.extend([0] * (max_len - len(current_parts)))
            required_parts.extend([0] * (max_len - len(required_parts)))
            return current_parts >= required_parts
        except Exception:
            return True

    def check_all_packages(self) -> bool:
        print_status("开始检查环境依赖...", "info")

        for import_name, pip_name in self.required_packages.items():
            if not self.check_package_installed(import_name):
                self.missing_packages.append((import_name, pip_name))
            elif import_name in self.version_requirements:
                current_version = self.get_package_version(import_name)
                required_version = self.version_requirements[import_name]
                if current_version and not self.compare_version(current_version, required_version):
                    print_status(
                        f"{import_name} 版本过低 (当前: {current_version}, 需要: >={required_version})",
                        "warning",
                    )
                    self.packages_need_upgrade.append((import_name, pip_name))

        for package_name, install_url in self.special_packages.items():
            if not self.check_package_installed(package_name):
                self.missing_packages.append((package_name, install_url))

        all_ok = not self.missing_packages and not self.packages_need_upgrade

        if all_ok:
            print_status("✓ 所有依赖包检查完成，环境正常", "success")
            return True

        if self.missing_packages:
            print_status(f"发现 {len(self.missing_packages)} 个缺失的包", "warning")
        if self.packages_need_upgrade:
            print_status(f"发现 {len(self.packages_need_upgrade)} 个需要升级的包", "warning")

        return False

    def install_missing_packages(self, auto_install: bool = True) -> bool:
        if not self.missing_packages and not self.packages_need_upgrade:
            return True

        if not auto_install:
            if self.missing_packages:
                print_status("缺失以下包:", "warning")
                for import_name, pip_name in self.missing_packages:
                    print_status(f"  - {import_name} ({pip_name})", "warning")
            if self.packages_need_upgrade:
                print_status("需要升级以下包:", "warning")
                for import_name, pip_name in self.packages_need_upgrade:
                    print_status(f"  - {import_name} ({pip_name})", "warning")
            return False

        if self.missing_packages:
            pkgs = [pip_name for _, pip_name in self.missing_packages]
            if not _install_packages(pkgs, label="unilabos"):
                self.failed_installs.extend(self.missing_packages)

        if self.packages_need_upgrade:
            pkgs = [pip_name for _, pip_name in self.packages_need_upgrade]
            if not _install_packages(pkgs, upgrade=True, label="unilabos"):
                self.failed_installs.extend(self.packages_need_upgrade)

        return not self.failed_installs

    def verify_installation(self) -> bool:
        if not self.missing_packages and not self.packages_need_upgrade:
            return True

        print_status("验证安装结果...", "info")
        failed_verification = []

        for import_name, pip_name in self.missing_packages:
            if not self.check_package_installed(import_name):
                failed_verification.append((import_name, pip_name))

        for import_name, pip_name in self.packages_need_upgrade:
            if not self.check_package_installed(import_name):
                failed_verification.append((import_name, pip_name))
            elif import_name in self.version_requirements:
                current_version = self.get_package_version(import_name)
                required_version = self.version_requirements[import_name]
                if current_version and not self.compare_version(current_version, required_version):
                    failed_verification.append((import_name, pip_name))
                    print_status(
                        f"  {import_name} 版本仍然过低 (当前: {current_version}, 需要: >={required_version})",
                        "error",
                    )

        if failed_verification:
            print_status(f"有 {len(failed_verification)} 个包验证失败:", "error")
            for import_name, pip_name in failed_verification:
                print_status(f"  - {import_name}", "error")
            return False

        print_status("✓ 所有包验证通过", "success")
        return True


def check_environment(auto_install: bool = True, show_details: bool = True) -> bool:
    """
    检查环境并自动安装缺失的包

    Args:
        auto_install: 是否自动安装缺失的包
        show_details: 是否显示详细信息

    Returns:
        bool: 环境检查是否通过
    """
    checker = EnvironmentChecker()

    if checker.check_all_packages():
        return True

    if not checker.install_missing_packages(auto_install):
        if show_details:
            print_status("请手动安装缺失的包后重新启动程序", "error")
        return False

    if not checker.verify_installation():
        if show_details:
            print_status("安装验证失败，请检查网络连接或手动安装", "error")
        return False

    return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="UniLabOS 环境依赖检查工具")
    parser.add_argument("--no-auto-install", action="store_true", help="仅检查环境，不自动安装缺失的包")
    parser.add_argument("--silent", action="store_true", help="静默模式，不显示详细信息")

    args = parser.parse_args()

    auto_install = not args.no_auto_install
    show_details = not args.silent

    success = check_environment(auto_install=auto_install, show_details=show_details)

    if not success:
        if show_details:
            print_status("环境检查失败", "error")
        sys.exit(1)
    else:
        if show_details:
            print_status("环境检查完成", "success")
