# /root/code/python/ieops-python/sdk/src/utils/subprocess_manager.py
"""
通用子进程管理器

提供子进程的启动、停止、监控、日志转发等功能，
可被各种 Server 复用（如 ComfyUIServer、vLLM Server 等）。

使用示例:
    manager = SubprocessManager(
        start_cmd="python main.py --port 8188",
        work_dir="/app/comfyui",
        health_check=lambda: check_http("http://localhost:8188/health"),
        startup_timeout=120,
        exit_on_failure=True,
    )
    
    # 启动子进程
    await manager.start()
    
    # 检查状态
    if manager.is_running:
        print("Process is running")
    
    # 停止子进程
    await manager.stop()
"""

import asyncio
import atexit
import os
import signal
import subprocess
import sys
import time
from typing import Any, Callable, Coroutine, Dict, List, Optional

from .log import uvicorn_logger

# 全局注册的子进程管理器列表，用于 atexit 清理
_registered_managers: List["SubprocessManager"] = []


def _cleanup_all_subprocesses():
    """atexit 回调：清理所有注册的子进程"""
    for manager in _registered_managers:
        manager._sync_kill()


# 注册 atexit 清理函数
atexit.register(_cleanup_all_subprocesses)


class SubprocessManager:
    """
    通用子进程管理器
    
    功能：
    - 启动子进程并等待就绪
    - 转发子进程日志
    - 后台监控子进程状态
    - 优雅停止子进程
    - 异常退出时可选择终止主进程
    """
    
    def __init__(
        self,
        start_cmd: str,
        work_dir: str,
        *,
        name: str = "subprocess",
        extra_env: Optional[Dict[str, str]] = None,
        health_check: Optional[Callable[[], Coroutine[Any, Any, bool]]] = None,
        startup_timeout: int = 120,
        health_check_interval: float = 1.0,
        monitor_interval: float = 5.0,
        graceful_shutdown_timeout: float = 5.0,
        exit_on_failure: bool = True,
        log_prefix: Optional[str] = None,
    ) -> None:
        """
        Args:
            start_cmd: 启动命令，如 "python main.py --port 8188"
            work_dir: 工作目录
            name: 进程名称（用于日志）
            extra_env: 额外的环境变量
            health_check: 健康检查函数，返回 True 表示就绪
            startup_timeout: 启动超时时间（秒）
            health_check_interval: 健康检查间隔（秒）
            monitor_interval: 后台监控间隔（秒）
            graceful_shutdown_timeout: 优雅关闭超时（秒）
            exit_on_failure: 子进程异常退出时是否调用 sys.exit(1)
            log_prefix: 日志前缀，默认使用 name
        """
        self._start_cmd = start_cmd
        self._work_dir = work_dir
        self._name = name
        self._extra_env = extra_env or {}
        self._health_check = health_check
        self._startup_timeout = startup_timeout
        self._health_check_interval = health_check_interval
        self._monitor_interval = monitor_interval
        self._graceful_shutdown_timeout = graceful_shutdown_timeout
        self._exit_on_failure = exit_on_failure
        self._log_prefix = log_prefix or f"[{name}]"
        
        self._process: Optional[subprocess.Popen] = None
        self._is_ready = False
        self._shutting_down = False
        self._monitor_task: Optional[asyncio.Task] = None
        self._log_task: Optional[asyncio.Task] = None

    @property
    def is_running(self) -> bool:
        """检查子进程是否正在运行"""
        return self._process is not None and self._process.poll() is None

    @property
    def is_ready(self) -> bool:
        """检查子进程是否已就绪（通过健康检查）"""
        return self._is_ready

    @property
    def exit_code(self) -> Optional[int]:
        """获取子进程退出码，如果还在运行返回 None"""
        if self._process is None:
            return None
        return self._process.poll()

    @property
    def pid(self) -> Optional[int]:
        """获取子进程 PID"""
        if self._process is None:
            return None
        return self._process.pid

    async def start(self) -> None:
        """
        启动子进程
        
        1. 启动进程
        2. 启动日志转发
        3. 等待健康检查通过
        4. 启动后台监控
        
        Raises:
            RuntimeError: 进程启动失败或在启动过程中退出
            TimeoutError: 启动超时
        """
        uvicorn_logger.info(f"{self._log_prefix} Starting: {self._start_cmd}")
        uvicorn_logger.info(f"{self._log_prefix} Work directory: {self._work_dir}")
        
        # 准备环境变量
        env = os.environ.copy()
        env.update(self._extra_env)
        
        # 启动子进程
        self._process = subprocess.Popen(
            self._start_cmd,
            shell=True,
            cwd=self._work_dir,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            preexec_fn=self._setup_child_process,  # 设置子进程属性
        )
        
        uvicorn_logger.info(f"{self._log_prefix} Started with PID: {self._process.pid}")
        
        # 注册到全局管理器列表，用于 atexit 清理
        _registered_managers.append(self)
        
        # 启动日志转发
        self._log_task = asyncio.create_task(self._forward_logs())
        
        # 等待健康检查通过
        await self._wait_for_ready()
        
        # 启动后台监控
        self._monitor_task = asyncio.create_task(self._monitor_process())
    
    def _setup_child_process(self) -> None:
        """
        子进程启动前的设置（在 fork 后、exec 前调用）
        
        1. 创建新的进程组（方便后续 kill 整个组）
        2. 设置 prctl PR_SET_PDEATHSIG（Linux 特性：父进程死亡时子进程收到信号）
        """
        # 创建新的进程组
        os.setsid()
        
        # Linux 特性：设置父进程死亡时子进程收到的信号
        try:
            import ctypes
            libc = ctypes.CDLL("libc.so.6", use_errno=True)
            PR_SET_PDEATHSIG = 1
            # 当父进程死亡时，子进程收到 SIGKILL
            libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL)
        except Exception:
            pass  # 非 Linux 系统或没有 prctl 支持，忽略

    async def stop(self) -> None:
        """
        优雅停止子进程
        
        1. 发送 SIGTERM
        2. 等待进程退出
        3. 超时后发送 SIGKILL
        """
        self._shutting_down = True
        
        # 取消监控任务
        if self._monitor_task:
            self._monitor_task.cancel()
            try:
                await self._monitor_task
            except asyncio.CancelledError:
                pass
        
        if not self._process:
            return
        
        uvicorn_logger.info(f"{self._log_prefix} Stopping process...")
        
        try:
            # 先尝试 SIGTERM
            os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
            
            # 等待进程退出
            wait_count = int(self._graceful_shutdown_timeout / 0.5)
            for _ in range(wait_count):
                if self._process.poll() is not None:
                    uvicorn_logger.info(f"{self._log_prefix} Process stopped gracefully")
                    return
                await asyncio.sleep(0.5)
            
            # 强制 SIGKILL
            uvicorn_logger.warning(f"{self._log_prefix} Force killing process...")
            os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)
            self._process.wait()
            uvicorn_logger.info(f"{self._log_prefix} Process killed")
            
        except ProcessLookupError:
            pass  # 进程已经退出
        except Exception as e:
            uvicorn_logger.error(f"{self._log_prefix} Error stopping process: {e}")
        finally:
            self._process = None
            self._is_ready = False
            # 从全局列表中移除
            if self in _registered_managers:
                _registered_managers.remove(self)
    
    def _sync_kill(self) -> None:
        """
        同步强制杀死子进程（用于 atexit 等无法使用 async 的场景）
        """
        if not self._process:
            return
        
        try:
            # 直接发送 SIGKILL 到整个进程组
            os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)
            self._process.wait(timeout=2)
        except ProcessLookupError:
            pass  # 进程已经退出
        except Exception:
            pass  # 忽略错误，这是清理阶段
        finally:
            self._process = None

    async def _forward_logs(self) -> None:
        """转发子进程的日志输出"""
        if not self._process or not self._process.stdout:
            return
        
        loop = asyncio.get_event_loop()
        
        def read_line():
            if self._process and self._process.stdout:
                return self._process.stdout.readline()
            return b""
        
        while self._process and self._process.poll() is None:
            try:
                line = await loop.run_in_executor(None, read_line)
                if line:
                    log_line = line.decode("utf-8", errors="ignore").strip()
                    uvicorn_logger.info(f"{self._log_prefix} {log_line}")
            except Exception:
                break

    async def _wait_for_ready(self) -> None:
        """等待子进程就绪"""
        if not self._health_check:
            # 没有健康检查，直接认为就绪
            self._is_ready = True
            uvicorn_logger.info(f"{self._log_prefix} No health check configured, assuming ready")
            return
        
        uvicorn_logger.info(f"{self._log_prefix} Waiting for process to be ready...")
        
        start_time = time.time()
        
        while time.time() - start_time < self._startup_timeout:
            try:
                if await self._health_check():
                    self._is_ready = True
                    uvicorn_logger.info(f"{self._log_prefix} Process is ready!")
                    return
            except Exception:
                pass
            
            # 检查进程是否还活着
            if self._process and self._process.poll() is not None:
                exit_code = self._process.returncode
                uvicorn_logger.error(
                    f"{self._log_prefix} Process died during startup with exit code {exit_code}"
                )
                if self._exit_on_failure:
                    sys.exit(1)
                raise RuntimeError(f"Process died with exit code {exit_code}")
            
            await asyncio.sleep(self._health_check_interval)
        
        # 启动超时
        uvicorn_logger.error(
            f"{self._log_prefix} Process failed to start within {self._startup_timeout}s"
        )
        if self._exit_on_failure:
            sys.exit(1)
        raise TimeoutError(f"Process failed to start within {self._startup_timeout}s")

    async def _monitor_process(self) -> None:
        """后台监控子进程状态，异常退出时处理"""
        while True:
            await asyncio.sleep(self._monitor_interval)
            
            # 如果正在正常关闭，停止监控
            if self._shutting_down:
                break
            
            if self._process is None:
                break
            
            exit_code = self._process.poll()
            if exit_code is not None:
                # 进程已退出且不是正常关闭
                uvicorn_logger.error(
                    f"{self._log_prefix} Process exited unexpectedly with code {exit_code}"
                )
                if self._exit_on_failure:
                    uvicorn_logger.error(f"{self._log_prefix} Shutting down main process...")
                    sys.exit(1)
                break


async def http_health_check(
    url: str,
    timeout: float = 5.0,
    expected_status: int = 200
) -> bool:
    """
    通用 HTTP 健康检查函数
    
    Args:
        url: 健康检查 URL
        timeout: 超时时间（秒）
        expected_status: 期望的状态码
    
    Returns:
        True 如果健康检查通过
    """
    import aiohttp
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(
                url,
                timeout=aiohttp.ClientTimeout(total=timeout)
            ) as resp:
                return resp.status == expected_status
    except Exception:
        return False


def create_http_health_check(
    url: str,
    timeout: float = 5.0,
    expected_status: int = 200
) -> Callable[[], Coroutine[Any, Any, bool]]:
    """
    创建 HTTP 健康检查函数的工厂方法
    
    使用示例:
        health_check = create_http_health_check("http://localhost:8188/system_stats")
        manager = SubprocessManager(
            start_cmd="...",
            work_dir="...",
            health_check=health_check,
        )
    """
    async def check() -> bool:
        return await http_health_check(url, timeout, expected_status)
    return check

