"""虚拟环境管理模块

使用 uv 进行依赖管理，提供极速安装和智能冲突解决能力。
"""

import hashlib
import os
import subprocess
import time
import fcntl
from pathlib import Path
from typing import List, Optional
from contextlib import contextmanager

from .config import PlaywrightExecutorConfig
from .exceptions import DependencyInstallError, RequirementNotFoundError
from .logger import logger


class EnvManager:
    """虚拟环境管理器"""

    def __init__(self, config: PlaywrightExecutorConfig):
        self.config = config
        self.venv_path = config.get_venv_path()
        self.service_path = config.get_service_path()
        
        # 并发控制
        self._lock_dir = Path("/workspace/venvs/.locks")
        self._ensure_lock_dir()
        
        # 启动时清理可能的残留锁文件
        self._cleanup_startup_locks()

    def _ensure_lock_dir(self) -> None:
        """确保锁文件目录存在"""
        try:
            self._lock_dir.mkdir(parents=True, exist_ok=True)
            logger.LOCAL.debug(f"锁文件目录已准备: {self._lock_dir}")
        except Exception as e:
            logger.LOCAL.warning(f"创建锁文件目录失败: {e}")

    def _cleanup_startup_locks(self) -> None:
        """启动时清理所有可能的残留锁文件"""
        try:
            # 清理环境锁目录中的所有锁文件
            if self._lock_dir.exists():
                for lock_file in self._lock_dir.glob("*.lock"):
                    try:
                        venv_name = lock_file.stem
                        self._cleanup_orphaned_locks(venv_name)
                    except Exception as e:
                        logger.LOCAL.warning(f"[启动清理] 清理环境锁失败 {lock_file}: {e}")
            
            # 清理uv相关锁文件
            self._cleanup_lock_files()
            logger.LOCAL.debug("[启动清理] 残留锁文件清理完成")
            
        except Exception as e:
            logger.LOCAL.warning(f"[启动清理] 启动锁清理异常: {e}")

    @contextmanager
    def _venv_lock(self):
        """虚拟环境操作的文件系统锁（增强异常安全）"""
        venv_name = self.config.get_venv_name()
        lock_file_path = self._lock_dir / f"{venv_name}.lock"
        lock_file = None
        
        logger.LOCAL.info(f"[并发控制] 请求环境锁: {venv_name}")
        
        try:
            # 清理可能的孤立锁文件
            self._cleanup_orphaned_locks(venv_name)
            
            # 创建锁文件并写入当前进程信息
            lock_file = open(lock_file_path, 'w')
            lock_file.write(f"{os.getpid()}\n{time.time()}\n")
            lock_file.flush()
            
            # 获取独占锁（阻塞直到获得）
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            logger.LOCAL.info(f"[并发控制] 已获得环境锁: {venv_name} (PID: {os.getpid()})")
            
            try:
                yield
            except Exception as e:
                logger.REMOTE.error(f"[并发控制] 锁保护的操作失败: {e}")
                raise
            finally:
                # 在任何情况下都要尝试释放锁
                logger.LOCAL.info(f"[并发控制] 释放环境锁: {venv_name}")
                
        except (IOError, OSError) as e:
            logger.REMOTE.error(f"[并发控制] 锁文件操作失败: {e}")
            raise DependencyInstallError(f"无法获得环境锁: {e}")
        except Exception as e:
            logger.REMOTE.error(f"[并发控制] 锁操作异常: {e}")
            raise
        finally:
            # 确保文件被关闭和删除
            if lock_file:
                try:
                    lock_file.close()
                except:
                    pass
            
            # 尝试删除锁文件（清理残留）
            try:
                if lock_file_path.exists():
                    lock_file_path.unlink()
                    logger.LOCAL.debug(f"[并发控制] 锁文件已清理: {lock_file_path}")
            except Exception as e:
                logger.LOCAL.warning(f"[并发控制] 锁文件清理失败: {e}")

    def _cleanup_orphaned_locks(self, venv_name: str) -> None:
        """清理孤立的锁文件（进程已结束但锁文件残留）"""
        lock_file_path = self._lock_dir / f"{venv_name}.lock"
        
        if not lock_file_path.exists():
            return
        
        try:
            # 读取锁文件内容
            content = lock_file_path.read_text().strip().split('\n')
            if len(content) < 2:
                logger.LOCAL.warning(f"[并发控制] 锁文件格式异常，强制清理: {lock_file_path}")
                lock_file_path.unlink()
                return
            
            pid = int(content[0])
            timestamp = float(content[1])
            
            # 检查进程是否还存在
            if self._is_process_alive(pid):
                # 检查锁文件是否过期（超过30分钟认为异常）
                if time.time() - timestamp > 1800:  # 30分钟
                    logger.LOCAL.warning(f"[并发控制] 锁文件超时，可能存在死锁，强制清理: PID={pid}")
                    lock_file_path.unlink()
                else:
                    logger.LOCAL.debug(f"[并发控制] 检测到有效锁文件，等待进程 PID={pid}")
            else:
                logger.LOCAL.warning(f"[并发控制] 检测到孤立锁文件，进程已结束 PID={pid}，清理中...")
                lock_file_path.unlink()
                
        except (ValueError, OSError) as e:
            logger.LOCAL.warning(f"[并发控制] 锁文件检查失败，强制清理: {e}")
            try:
                lock_file_path.unlink()
            except:
                pass
        except Exception as e:
            logger.REMOTE.error(f"[并发控制] 孤立锁清理异常: {e}")

    def _is_process_alive(self, pid: int) -> bool:
        """检查指定PID的进程是否还存在"""
        try:
            # 发送信号0检查进程是否存在（不会杀死进程）
            os.kill(pid, 0)
            return True
        except OSError:
            return False
        except Exception:
            return False

    def ensure_environment(self) -> None:
        """
        确保虚拟环境存在且依赖已安装（并发安全）

        流程:
            1. 验证业务逻辑文件夹和 main.py 存在
            2. 获取文件系统锁（防止并发冲突）
            3. 创建虚拟环境（如果不存在）
            4. 检查依赖是否需要更新（增强MD5校验）
            5. 安装/更新依赖
        """
        # 1. 验证业务逻辑文件夹
        self._validate_service_folder()

        # 2. 使用文件系统锁保护并发操作
        with self._venv_lock():
            # 3. 双重检查：获得锁后再次检查环境状态
            if self._is_environment_ready():
                logger.LOCAL.info(f"[并发控制] 环境已就绪，跳过创建: {self.venv_path}")
                return

            # 4. 创建虚拟环境
            if not self.venv_path.exists():
                logger.LOCAL.info(f"虚拟环境不存在，开始创建: {self.venv_path}")
                self._create_venv()
            else:
                logger.LOCAL.info(f"虚拟环境已存在: {self.venv_path}")

            # 5. 检查并安装依赖
            self._ensure_dependencies()

    def _validate_service_folder(self) -> None:
        """验证业务逻辑文件夹和 main.py 是否存在"""
        if not self.service_path.exists():
            raise RequirementNotFoundError(f"业务逻辑文件夹不存在: {self.service_path}")

        main_py = self.service_path / "main.py"
        if not main_py.exists():
            raise RequirementNotFoundError(f"main.py 不存在: {main_py}")

        logger.LOCAL.info(f"业务逻辑文件夹验证通过: {self.service_path}")

    def _is_environment_ready(self) -> bool:
        """
        检查环境是否完全就绪（虚拟环境 + 依赖）
        
        用于双重检查模式：获得锁后快速判断是否可以跳过所有操作
        
        Returns:
            True: 环境完全就绪，可以跳过所有操作
            False: 需要进行环境创建或依赖安装
        """
        try:
            # 1. 检查虚拟环境是否存在且有效
            if not self._is_venv_valid():
                logger.LOCAL.debug("[环境检查] 虚拟环境无效")
                return False
            
            # 2. 检查依赖是否为最新（MD5校验）
            requirements_files = self._get_requirements_files()
            if not requirements_files:
                logger.LOCAL.debug("[环境检查] 无依赖文件，环境就绪")
                return True
            
            current_md5 = self._calculate_requirements_md5(requirements_files)
            cached_md5 = self._get_cached_md5()
            
            if current_md5 == cached_md5:
                logger.LOCAL.info("[环境检查] 依赖文件未变化且虚拟环境有效，环境完全就绪")
                return True
            else:
                logger.LOCAL.debug(f"[环境检查] 依赖文件有变化: {cached_md5} -> {current_md5}")
                return False
                
        except Exception as e:
            logger.LOCAL.warning(f"[环境检查] 检查过程异常，将执行完整流程: {e}")
            return False

    def _is_venv_valid(self) -> bool:
        """检查虚拟环境是否存在且有效"""
        if not self.venv_path.exists():
            return False
        
        # 检查关键文件是否存在
        python_path = self.venv_path / "bin" / "python"
        if not python_path.exists():
            logger.LOCAL.warning(f"虚拟环境Python解释器不存在: {python_path}")
            return False
        
        # 检查 site-packages 目录
        site_packages = self.get_site_packages_paths()
        if not site_packages:
            logger.LOCAL.warning("虚拟环境site-packages目录不存在")
            return False
            
        return True

    def _create_venv(self) -> None:
        """创建虚拟环境（使用 uv）"""
        try:
            self.venv_path.parent.mkdir(parents=True, exist_ok=True)

            # 使用 uv 创建虚拟环境（速度更快）
            cmd = ["uv", "venv", str(self.venv_path), "--python", "3.12"]
            logger.LOCAL.info(f"执行命令: {' '.join(cmd)}")

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)

            if result.returncode != 0:
                raise DependencyInstallError(f"创建虚拟环境失败: {result.stderr}")

            logger.LOCAL.info("虚拟环境创建成功（使用 uv）")

        except subprocess.TimeoutExpired:
            raise DependencyInstallError("创建虚拟环境超时（60秒）")
        except Exception as e:
            raise DependencyInstallError(f"创建虚拟环境异常: {str(e)}")

    def _ensure_dependencies(self) -> None:
        """确保依赖已安装且是最新的（增强MD5校验，并发安全）"""
        requirements_files = self._get_requirements_files()

        if not requirements_files:
            logger.LOCAL.info("[依赖管理] 未找到 pyproject.toml 文件，跳过依赖安装")
            return

        # 计算当前 requirements 文件的 MD5
        current_md5 = self._calculate_requirements_md5(requirements_files)
        logger.LOCAL.debug(f"[依赖管理] 当前依赖MD5: {current_md5}")

        # 获取缓存的 MD5
        cached_md5 = self._get_cached_md5()
        logger.LOCAL.debug(f"[依赖管理] 缓存依赖MD5: {cached_md5}")

        if current_md5 == cached_md5:
            logger.LOCAL.info("[依赖管理] 依赖文件未变化（MD5一致），跳过安装")
            return

        # 安装依赖
        logger.LOCAL.info(f"[依赖管理] 依赖文件有变化，开始安装... ({len(requirements_files)}个配置文件)")
        
        install_success = True
        for i, req_file in enumerate(requirements_files, 1):
            try:
                logger.LOCAL.info(f"[依赖管理] 安装进度 {i}/{len(requirements_files)}: {req_file.name}")
                self._install_requirements(req_file)
            except Exception as e:
                install_success = False
                logger.REMOTE.error(f"[依赖管理] 安装失败 {req_file.name}: {e}")
                raise

        # 更新 MD5 缓存（仅在全部安装成功后）
        if install_success:
            self._save_md5_cache(current_md5)
            logger.LOCAL.info("[依赖管理] 所有依赖安装完成，MD5缓存已更新")
        else:
            logger.REMOTE.error("[依赖管理] 依赖安装失败，未更新MD5缓存")

    def _get_requirements_files(self) -> List[Path]:
        """
        获取需要安装的依赖配置文件列表

        优先级:
            1. 根目录的 pyproject.toml（通用依赖）
            2. 业务逻辑文件夹的 pyproject.toml（业务特定依赖）
        """
        files = []

        # 1. 根目录的 pyproject.toml
        root_pyproject = self.config.git_repo_dir / "pyproject.toml"
        if root_pyproject.exists():
            files.append(root_pyproject)
            logger.LOCAL.info(f"发现根目录依赖文件: {root_pyproject}")

        # 2. 业务逻辑文件夹的 pyproject.toml
        service_pyproject = self.service_path / "pyproject.toml"
        if service_pyproject.exists():
            files.append(service_pyproject)
            logger.LOCAL.info(f"发现业务目录依赖文件: {service_pyproject}")

        return files

    def _calculate_requirements_md5(self, files: List[Path]) -> str:
        """计算多个依赖配置文件的联合 MD5"""
        md5_hash = hashlib.md5()

        for file in sorted(files, key=lambda x: str(x)):  # 排序确保顺序一致
            with open(file, "rb") as f:
                md5_hash.update(f.read())

        return md5_hash.hexdigest()

    def _get_cached_md5(self) -> Optional[str]:
        """获取缓存的 MD5 值"""
        md5_file = self._get_md5_cache_file()

        if not md5_file.exists():
            return None

        try:
            return md5_file.read_text().strip()
        except Exception:
            return None

    def _save_md5_cache(self, md5: str) -> None:
        """保存 MD5 缓存"""
        md5_file = self._get_md5_cache_file()
        md5_file.parent.mkdir(parents=True, exist_ok=True)
        md5_file.write_text(md5)

    def _get_md5_cache_file(self) -> Path:
        """获取 MD5 缓存文件路径"""
        venv_name = self.config.get_venv_name()
        return self.config.md5_cache_dir / f"{venv_name}.md5"

    def _cleanup_lock_files(self) -> None:
        """清理可能的 uv 锁文件（增强异常处理）"""
        cleanup_targets = [
            (self.venv_path / ".lock", "虚拟环境锁文件"),
            (Path.home() / ".cache" / "uv" / ".lock", "uv全局缓存锁"),
            (Path("/tmp") / "uv.lock", "uv临时锁文件"),  # 额外的可能锁位置
        ]
        
        for lock_file, description in cleanup_targets:
            try:
                if lock_file.exists():
                    # 检查锁文件是否可以安全删除
                    if self._can_safely_remove_lock(lock_file):
                        lock_file.unlink()
                        logger.LOCAL.info(f"[锁清理] 已清理{description}: {lock_file}")
                    else:
                        logger.LOCAL.warning(f"[锁清理] {description}可能被其他进程使用，跳过: {lock_file}")
            except (OSError, IOError) as e:
                logger.LOCAL.warning(f"[锁清理] 清理{description}失败: {e}")
            except Exception as e:
                logger.REMOTE.error(f"[锁清理] 清理{description}异常: {e}")

    def _can_safely_remove_lock(self, lock_file: Path) -> bool:
        """检查锁文件是否可以安全删除"""
        try:
            # 尝试以非阻塞方式获得锁，如果成功说明没有其他进程在使用
            with open(lock_file, 'r') as f:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
                    return True  # 成功获得锁，说明可以安全删除
                except BlockingIOError:
                    return False  # 锁被占用，不能删除
        except Exception:
            return True  # 如果无法检查，默认允许删除

    def _install_requirements(self, config_file: Path) -> None:
        """安装指定的依赖配置文件（使用 uv）"""
        try:
            python_path = self.venv_path / "bin" / "python"
            
            # 清理可能的锁文件
            self._cleanup_lock_files()
            
            cmd = [
                "uv", "pip", "install",
                "--python", str(python_path),
                "-e", str(config_file.parent),
                "--index-url", "https://pypi.tuna.tsinghua.edu.cn/simple",  # 国内镜像源
                "--extra-index-url", "https://pypi.org/simple",  # 备用源
            ]
            
            logger.LOCAL.info(f"安装依赖: {config_file}")
            logger.LOCAL.info(f"执行命令: {' '.join(cmd)}")
            logger.LOCAL.info(f"配置文件内容: {config_file.read_text()[:500]}...")
            
            # 设置环境变量（Dockerfile已预配置大部分uv优化项）
            env = os.environ.copy()
            env["PYTHONUNBUFFERED"] = "1"  # Python 不缓冲
            
            logger.LOCAL.info("⏳ 正在安装依赖，请耐心等待...")
            
            try:
                result = subprocess.run(
                    cmd,
                    env=env,
                    capture_output=True,
                    text=True,
                    timeout=600  # 10分钟总超时（容器环境）
                )
                
                # 输出所有日志
                if result.stdout:
                    for line in result.stdout.splitlines():
                        line = line.strip()
                        if line:
                            logger.LOCAL.info(f"[uv] {line}")
                
                if result.returncode != 0:
                    error_msg = result.stderr or result.stdout or "无输出"
                    raise DependencyInstallError(f"安装依赖失败({config_file}): {error_msg[-1000:]}")
                
                logger.LOCAL.info(f"✅ 依赖安装成功: {config_file}")
                
            except subprocess.TimeoutExpired as e:
                # 超时后尝试读取部分输出
                stdout = e.stdout.decode() if e.stdout else ""
                stderr = e.stderr.decode() if e.stderr else ""
                output = stdout or stderr or "无输出"
                
                raise DependencyInstallError(
                    f"安装依赖超时({config_file})，超过600秒\n"
                    f"建议检查网络环境或镜像源配置\n"
                    f"部分输出:\n{output[-1000:]}"
                )
            
        except DependencyInstallError:
            raise
        except Exception as e:
            raise DependencyInstallError(f"安装依赖异常({config_file}): {str(e)}")

    def get_site_packages_paths(self) -> List[str]:
        """
        获取虚拟环境的 site-packages 路径

        用于动态导入时添加到 sys.path

        Returns:
            存在的 site-packages 路径列表
        """
        python_version = self.config.get_python_version()

        site_packages_paths = [
            str(self.venv_path / "lib" / python_version / "site-packages"),
            str(self.venv_path / "lib64" / python_version / "site-packages"),  # 兼容某些系统
        ]

        # 只返回存在的路径
        existing_paths = [p for p in site_packages_paths if Path(p).exists()]

        if not existing_paths:
            logger.LOCAL.warning(f"未找到 site-packages 目录: {site_packages_paths}")

        return existing_paths