"""
生成阶段状态管理器
用于支持生成任务的暂停、停止功能
使用轮询检查机制，确保线程安全
支持多用户/多generation_id并行管理
"""

import threading
import asyncio
import time
from typing import Optional, Dict
from enum import Enum


class GenerationStatus(Enum):
    """生成状态枚举"""
    GENERATING = "generating"  # 正在生成
    PAUSED = "paused"          # 已暂停
    STOPPED = "stopped"        # 已停止
    COMPLETED = "completed"    # 已完成
    IDLE = "idle"              # 空闲


class GenerationStateManager:
    """
    生成阶段状态管理器

    功能：
    - 管理多用户/多generation_id的并行状态（使用字典存储）
    - 管理生成任务状态（生成中/暂停/停止/完成）
    - 提供线程安全的状态查询和修改接口

    设计说明：
    - 保持单例模式（Manager本身）
    - 但内部状态按 generation_id 隔离存储
    - 解决多用户并发生成时的干扰问题
    """

    _instance: Optional['GenerationStateManager'] = None
    _lock = threading.Lock()

    def __new__(cls):
        """单例模式"""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        """初始化状态管理器"""
        if self._initialized:
            return

        self._initialized = True

        # 状态存储：generation_id -> 状态字典
        # 结构：{
        #   'status': GenerationStatus,
        #   'goal': str,
        #   'should_stop': bool
        # }
        self._states: Dict[str, Dict] = {}

        # 每个 generation_id 的锁（更细粒度的锁）
        self._locks: Dict[str, threading.Lock] = {}

        # 全局锁（用于管理 _states 和 _locks 本身的线程安全）
        self._manager_lock = threading.Lock()

    def _get_lock(self, generation_id: str) -> threading.Lock:
        """获取指定 generation_id 的锁，如果不存在则创建"""
        with self._manager_lock:
            if generation_id not in self._locks:
                self._locks[generation_id] = threading.Lock()
            return self._locks[generation_id]

    def _ensure_state(self, generation_id: str, goal: str = None) -> Dict:
        """确保指定 generation_id 的状态存在"""
        with self._manager_lock:
            if generation_id not in self._states:
                self._states[generation_id] = {
                    'status': GenerationStatus.IDLE,
                    'goal': goal,
                    'should_stop': False
                }
            return self._states[generation_id]

    def _get_state(self, generation_id: str) -> Optional[Dict]:
        """获取指定 generation_id 的状态，不存在则返回 None"""
        with self._manager_lock:
            return self._states.get(generation_id)

    def _cleanup_state(self, generation_id: str):
        """清理指定 generation_id 的状态"""
        with self._manager_lock:
            self._states.pop(generation_id, None)
            self._locks.pop(generation_id, None)

    def get_status(self, generation_id: str) -> Optional[GenerationStatus]:
        """获取当前生成状态"""
        state = self._get_state(generation_id)
        if state is None:
            return None
        with self._get_lock(generation_id):
            return state['status']

    def start_generation(self, generation_id: str, goal: str):
        """开始生成"""
        state = self._ensure_state(generation_id, goal)
        with self._get_lock(generation_id):
            state['status'] = GenerationStatus.GENERATING
            state['goal'] = goal
            state['should_stop'] = False

    def stop_generation(self, generation_id: str) -> bool:
        """
        停止生成

        Args:
            generation_id: 生成ID

        Returns:
            bool: 是否成功停止（COMPLETED 状态也返回 True，表示已停止）
        """
        state = self._get_state(generation_id)
        if state is None:
            return True  # 不存在也算停止成功

        with self._get_lock(generation_id):
            if state['status'] == GenerationStatus.STOPPED:
                return True  # 已经停止也算成功

            if state['status'] == GenerationStatus.COMPLETED:
                return True  # 已完成也视为停止成功

            if state['status'] == GenerationStatus.IDLE:
                return True  # 空闲状态也视为无需停止

            # 真正需要停止的情况
            state['status'] = GenerationStatus.STOPPED
            state['should_stop'] = True
            return True

    def complete_generation(self, generation_id: str):
        """标记生成完成"""
        state = self._ensure_state(generation_id)
        with self._get_lock(generation_id):
            state['status'] = GenerationStatus.COMPLETED

    def cleanup(self, generation_id: str):
        """清理指定 generation_id 的所有状态"""
        self._cleanup_state(generation_id)

    def should_stop(self, generation_id: str) -> bool:
        """检查是否应该停止"""
        state = self._get_state(generation_id)
        if state is None:
            return False
        with self._get_lock(generation_id):
            return state.get('should_stop', False)

    def is_stopped(self, generation_id: str) -> bool:
        """检查是否已停止"""
        state = self._get_state(generation_id)
        if state is None:
            return False
        with self._get_lock(generation_id):
            return state['status'] == GenerationStatus.STOPPED

    def is_completed(self, generation_id: str) -> bool:
        """检查是否已完成"""
        state = self._get_state(generation_id)
        if state is None:
            return False
        with self._get_lock(generation_id):
            return state['status'] == GenerationStatus.COMPLETED

    def is_active(self, generation_id: str) -> bool:
        """检查是否处于活动状态（生成中或暂停中）"""
        state = self._get_state(generation_id)
        if state is None:
            return False
        with self._get_lock(generation_id):
            return state['status'] == GenerationStatus.GENERATING

    def check_and_set_stop(self, generation_id: str) -> bool:
        """
        检查是否应该停止，如果应该则设置停止状态

        Args:
            generation_id: 生成ID

        Returns:
            bool: True表示应该停止，False表示可以继续
        """
        state = self._get_state(generation_id)
        if state is None:
            return False
        with self._get_lock(generation_id):
            if state['should_stop']:
                return True
            return False

    def generate_id(self, goal: str) -> str:
        """
        生成唯一的 generation_id

        Args:
            goal: 生成目标

        Returns:
            str: 格式为 {goal}_{timestamp}
        """
        return f"{goal}_{int(time.time() * 1000)}"


# 全局单例实例
generation_state_manager = GenerationStateManager()
