#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
平滑重启模块
支持优雅停止、零停机重启、滚动更新
"""

import os
import sys
import signal
import time
import threading
import logging
import subprocess
from typing import Dict, List, Optional, Callable
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum

logger = logging.getLogger(__name__)


class RestartStrategy(Enum):
    """重启策略"""
    GRACEFUL = "graceful"  # 优雅停止，等待请求完成
    ROLLING = "rolling"     # 滚动更新，零停机
    IMMEDIATE = "immediate" # 立即重启


class ServerState(Enum):
    """服务器状态"""
    RUNNING = "running"
    STOPPING = "stopping"
    STOPPED = "stopped"
    STARTING = "starting"
    RESTARTING = "restarting"


@dataclass
class PendingRequest:
    """待处理的请求"""
    request_id: str
    start_time: float
    endpoint: str
    method: str
    client_address: tuple


class GracefulRestartManager:
    """平滑重启管理器"""

    def __init__(self, config: Dict = None):
        """
        初始化平滑重启管理器

        Args:
            config: 重启配置
        """
        self.config = config or {}
        self.state = ServerState.STOPPED

        # 等待超时时间
        self.graceful_timeout = self.config.get('graceful_timeout', 30)  # 秒
        self.shutdown_timeout = self.config.get('shutdown_timeout', 10)  # 秒

        # 待处理请求追踪
        self._pending_requests: Dict[str, PendingRequest] = {}
        self._requests_lock = threading.Lock()

        # 请求计数器
        self._request_counter = 0
        self._counter_lock = threading.Lock()

        # 回调函数
        self._on_prepare_restart: Optional[Callable] = None
        self._on_start_restart: Optional[Callable] = None
        self._on_complete_restart: Optional[Callable] = None
        self._on_restart_failed: Optional[Callable] = None

        # 状态锁
        self._state_lock = threading.Lock()

        # 重启历史
        self._restart_history: List[Dict] = []
        self._history_lock = threading.Lock()

        # 信号处理器
        self._setup_signal_handlers()

    def _setup_signal_handlers(self):
        """设置信号处理器"""
        signal.signal(signal.SIGTERM, self._handle_signal)
        signal.signal(signal.SIGINT, self._handle_signal)
        signal.signal(signal.SIGHUP, self._handle_signal)

    def _handle_signal(self, signum, frame):
        """处理退出信号"""
        signal_names = {
            signal.SIGTERM: 'SIGTERM',
            signal.SIGINT: 'SIGINT',
            signal.SIGHUP: 'SIGHUP'
        }
        sig_name = signal_names.get(signum, str(signum))
        logger.info(f"收到信号 {sig_name}，准备优雅关闭")

        # 设置状态为停止中
        self.state = ServerState.STOPPING

    def register_request(self, endpoint: str, method: str, client: tuple) -> str:
        """
        注册一个待处理的请求

        Args:
            endpoint: 请求端点
            method: 请求方法
            client: 客户端地址

        Returns:
            请求 ID
        """
        with self._counter_lock:
            self._request_counter += 1
            request_id = f"req_{self._request_counter}_{int(time.time())}"

        request = PendingRequest(
            request_id=request_id,
            start_time=time.time(),
            endpoint=endpoint,
            method=method,
            client_address=client
        )

        with self._requests_lock:
            self._pending_requests[request_id] = request

        return request_id

    def complete_request(self, request_id: str):
        """标记请求完成"""
        with self._requests_lock:
            self._pending_requests.pop(request_id, None)

    def get_pending_count(self) -> int:
        """获取待处理请求数量"""
        with self._requests_lock:
            return len(self._pending_requests)

    def get_pending_requests(self) -> List[Dict]:
        """获取待处理请求详情"""
        with self._requests_lock:
            return [
                {
                    'request_id': r.request_id,
                    'endpoint': r.endpoint,
                    'method': r.method,
                    'duration': round(time.time() - r.start_time, 2),
                    'client': str(r.client_address[0]) if r.client_address else 'unknown'
                }
                for r in self._pending_requests.values()
            ]

    def is_safe_to_restart(self) -> bool:
        """检查是否安全重启（没有待处理请求）"""
        return self.get_pending_count() == 0

    def wait_for_requests(
        self,
        timeout: int = None,
        progress_callback: Callable[[int, int], None] = None
    ) -> bool:
        """
        等待所有请求完成

        Args:
            timeout: 超时时间（秒）
            progress_callback: 进度回调函数(remaining, total)

        Returns:
            是否在超时前完成所有请求
        """
        timeout = timeout or self.graceful_timeout
        start_time = time.time()

        while time.time() - start_time < timeout:
            pending = self.get_pending_count()

            if progress_callback:
                progress_callback(pending, 0)

            if pending == 0:
                return True

            time.sleep(0.5)

        return False

    def set_prepare_restart_callback(self, callback: Callable):
        """设置准备重启回调"""
        self._on_prepare_restart = callback

    def set_start_restart_callback(self, callback: Callable):
        """设置开始重启回调"""
        self._on_start_restart = callback

    def set_complete_restart_callback(self, callback: Callable):
        """设置完成重启回调"""
        self._on_complete_restart = callback

    def set_restart_failed_callback(self, callback: Callable):
        """设置重启失败回调"""
        self._on_restart_failed = callback

    def prepare_restart(self) -> Dict:
        """
        准备重启（通知各模块准备）

        Returns:
            准备结果
        """
        with self._state_lock:
            self.state = ServerState.STOPPING

        result = {
            'success': True,
            'pending_requests': self.get_pending_count(),
            'message': ''
        }

        # 通知回调
        if self._on_prepare_restart:
            try:
                self._on_prepare_restart()
            except Exception as e:
                logger.error(f"Prepare restart callback failed: {e}")
                result['success'] = False
                result['message'] = str(e)

        return result

    def perform_restart(
        self,
        strategy: RestartStrategy = RestartStrategy.GRACEFUL,
        new_config: Dict = None,
        script_path: str = None
    ) -> Dict:
        """
        执行重启

        Args:
            strategy: 重启策略
            new_config: 新配置（用于配置热更新）
            script_path: 服务器脚本路径

        Returns:
            重启结果
        """
        with self._state_lock:
            if self.state == ServerState.RESTARTING:
                return {
                    'success': False,
                    'error': 'Restart already in progress'
                }
            self.state = ServerState.RESTARTING

        start_time = time.time()
        result = {
            'success': False,
            'strategy': strategy.value,
            'elapsed_seconds': 0,
            'message': ''
        }

        try:
            # 准备阶段
            prepare_result = self.prepare_restart()
            if not prepare_result['success']:
                result['message'] = f"Prepare failed: {prepare_result['message']}"
                self.state = ServerState.RUNNING
                return result

            # 根据策略执行重启
            if strategy == RestartStrategy.GRACEFUL:
                result = self._graceful_restart(prepare_result)
            elif strategy == RestartStrategy.ROLLING:
                result = self._rolling_restart(script_path, new_config)
            elif strategy == RestartStrategy.IMMEDIATE:
                result = self._immediate_restart(script_path)

            result['elapsed_seconds'] = round(time.time() - start_time, 2)

            # 记录重启历史
            self._add_to_history(result)

        except Exception as e:
            logger.error(f"Restart failed: {e}")
            result['success'] = False
            result['message'] = str(e)
            self.state = ServerState.RUNNING

            if self._on_restart_failed:
                try:
                    self._on_restart_failed(str(e))
                except Exception:
                    pass

        return result

    def _graceful_restart(self, prepare_result: Dict) -> Dict:
        """优雅重启（等待请求完成）"""
        result = {
            'success': True,
            'strategy': 'graceful',
            'pending_requests': prepare_result['pending_requests'],
            'message': ''
        }

        # 等待请求完成
        pending = prepare_result['pending_requests']
        if pending > 0:
            logger.info(f"等待 {pending} 个请求完成，超时 {self.graceful_timeout} 秒")

            def progress_callback(remaining, total):
                if remaining % 5 == 0:
                    logger.info(f"还有 {remaining} 个请求待处理")

            success = self.wait_for_requests(
                timeout=self.graceful_timeout,
                progress_callback=progress_callback
            )

            if not success:
                remaining = self.get_pending_count()
                result['success'] = False
                result['message'] = f"{remaining} 个请求未在 {self.graceful_timeout} 秒内完成"
                logger.warning(result['message'])
                self.state = ServerState.RUNNING
                return result

            result['message'] = '所有请求已完成'
            logger.info('所有请求已完成')

        # 通知开始重启
        if self._on_start_restart:
            try:
                self._on_start_restart()
            except Exception as e:
                logger.error(f"Start restart callback failed: {e}")

        result['success'] = True
        result['message'] = 'Graceful restart prepared'
        self.state = ServerState.STOPPED

        return result

    def _rolling_restart(self, script_path: str, new_config: Dict = None) -> Dict:
        """滚动重启（零停机）"""
        result = {
            'success': True,
            'strategy': 'rolling',
            'message': ''
        }

        if not script_path:
            script_path = os.path.join(os.path.dirname(__file__), '..', 'main.py')

        script_path = os.path.abspath(script_path)

        # 检查是否有新配置需要应用
        if new_config:
            result['config_updated'] = True
            logger.info("配置将在重启后应用")
        else:
            result['config_updated'] = False

        # 通知开始重启
        if self._on_start_restart:
            try:
                self._on_start_restart()
            except Exception as e:
                logger.error(f"Start restart callback failed: {e}")

        # 启动新进程（在同一端口，但使用不同进程ID）
        # 注意：实际实现需要在负载均衡器层面处理
        logger.info("滚动重启：建议在负载均衡器层面处理零停机")

        self.state = ServerState.STOPPED
        result['message'] = 'Rolling restart prepared - use load balancer for zero-downtime'

        return result

    def _immediate_restart(self, script_path: str = None) -> Dict:
        """立即重启"""
        result = {
            'success': True,
            'strategy': 'immediate',
            'message': ''
        }

        if not script_path:
            script_path = os.path.join(os.path.dirname(__file__), '..', 'main.py')

        script_path = os.path.abspath(script_path)

        # 通知开始重启
        if self._on_start_restart:
            try:
                self._on_start_restart()
            except Exception as e:
                logger.error(f"Start restart callback failed: {e}")

        # 发送重启信号给主进程
        logger.info("立即重启服务器")

        # 在子进程中重启
        try:
            # 启动新进程
            cmd = [sys.executable, script_path]
            env = os.environ.copy()
            env['HYC_RESTARTED'] = '1'

            subprocess.Popen(cmd, env=env)

            self.state = ServerState.STOPPED
            result['message'] = 'Immediate restart initiated'

        except Exception as e:
            result['success'] = False
            result['message'] = f"Failed to restart: {str(e)}"
            self.state = ServerState.RUNNING

        return result

    def _add_to_history(self, result: Dict):
        """添加重启历史记录"""
        record = {
            'timestamp': datetime.now().isoformat(),
            'strategy': result.get('strategy', 'unknown'),
            'success': result.get('success', False),
            'elapsed_seconds': result.get('elapsed_seconds', 0),
            'message': result.get('message', '')
        }

        with self._history_lock:
            self._restart_history.append(record)
            # 只保留最近 20 条记录
            self._restart_history = self._restart_history[-20:]

    def get_restart_history(self) -> List[Dict]:
        """获取重启历史"""
        with self._history_lock:
            return list(self._restart_history)

    def get_stats(self) -> Dict:
        """获取统计信息"""
        return {
            'state': self.state.value,
            'pending_requests': self.get_pending_count(),
            'graceful_timeout': self.graceful_timeout,
            'shutdown_timeout': self.shutdown_timeout,
            'restart_count': len(self._restart_history),
            'recent_restarts': self.get_restart_history()[-5:]
        }

    def update_config(self, new_config: Dict):
        """更新配置（热更新）"""
        if 'graceful_timeout' in new_config:
            self.graceful_timeout = new_config['graceful_timeout']
        if 'shutdown_timeout' in new_config:
            self.shutdown_timeout = new_config['shutdown_timeout']

        logger.info(f"重启配置已更新: timeout={self.graceful_timeout}s")


class ServerHealthChecker:
    """服务器健康检查器"""

    def __init__(self, host: str = 'localhost', port: int = 8080):
        self.host = host
        self.port = port
        self._last_check = None
        self._is_healthy = False

    def check(self) -> Dict:
        """
        检查服务器健康状态

        Returns:
            健康检查结果
        """
        import socket

        result = {
            'healthy': False,
            'latency_ms': 0,
            'error': None,
            'timestamp': datetime.now().isoformat()
        }

        try:
            start_time = time.time()

            # 尝试连接
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(5)
            sock.connect((self.host, self.port))
            sock.close()

            result['latency_ms'] = round((time.time() - start_time) * 1000, 2)
            result['healthy'] = True

        except Exception as e:
            result['error'] = str(e)

        self._last_check = result
        return result

    def get_latency(self) -> float:
        """获取延迟（毫秒）"""
        if self._last_check:
            return self._last_check.get('latency_ms', 0)
        return 0


class RollingRestartManager:
    """滚动重启管理器（支持多实例）"""

    def __init__(self, config: Dict = None):
        self.config = config or {}
        self.instances: Dict[str, Dict] = {}  # instance_id -> info
        self.current_instance_id = None
        self._lock = threading.Lock()

    def register_instance(self, instance_id: str, info: Dict):
        """注册实例"""
        with self._lock:
            self.instances[instance_id] = {
                **info,
                'registered_at': datetime.now().isoformat(),
                'status': 'active'
            }

    def unregister_instance(self, instance_id: str):
        """注销实例"""
        with self._lock:
            if instance_id in self.instances:
                self.instances[instance_id]['status'] = 'draining'

    def get_active_instances(self) -> List[str]:
        """获取活跃实例列表"""
        with self._lock:
            return [
                i for i, info in self.instances.items()
                if info['status'] == 'active'
            ]

    def perform_rolling_restart(
        self,
        instance_id: str,
        restart_func: Callable
    ) -> Dict:
        """
        对单个实例执行滚动重启

        Args:
            instance_id: 实例 ID
            restart_func: 重启函数

        Returns:
            重启结果
        """
        with self._lock:
            if instance_id not in self.instances:
                return {
                    'success': False,
                    'error': f'Instance {instance_id} not found'
                }

        # 标记为排水中
        self.instances[instance_id]['status'] = 'draining'

        # 等待连接耗尽
        time.sleep(5)

        # 执行重启
        try:
            restart_func(instance_id)

            self.instances[instance_id]['status'] = 'active'
            self.instances[instance_id]['restarted_at'] = datetime.now().isoformat()

            return {
                'success': True,
                'instance_id': instance_id
            }

        except Exception as e:
            self.instances[instance_id]['status'] = 'error'
            return {
                'success': False,
                'instance_id': instance_id,
                'error': str(e)
            }