import signal
import atexit
import asyncio as aio
import multiprocessing as mp
from typing import Any, Callable, Awaitable, Type, Coroutine

from setproctitle import setproctitle

from ldnode import logger
from .config import Settings
from .task_manager import TaskManager
from .events import Events
from .schemas import Worker, BaseContext, BaseWorkerState


class WorkerState(BaseWorkerState):
    """ 工作进程状态 """
    def __init__(self, worker_id: str, worker_cls: Type[Worker]):
        self.worker_id = worker_id
        self.worker_cls = worker_cls
        self.input = mp.Queue()
        self.output = mp.Queue()
        self.process: mp.Process | None = None
        self.send_data_cache: list[Any] = []
        self._used = True
        
    def get_worker_id(self) -> str:
        """ 获取工作进程ID """
        return self.worker_id
        
    def is_alive(self) -> bool:
        """ 检查工作进程是否存活 """
        return self.process is not None and self.process.is_alive()
    
    def is_used(self) -> bool:
        """ 检查工作进程是否被使用 """
        return self._used
    
    async def send(self, data: Any):
        """ 发送数据到工作进程 """
        if self.process and self.process.is_alive():
            await aio.to_thread(self.input.put, data)
        else:
            self.send_data_cache.append(data)
            
    async def recv(self) -> Any:
        """ 接收数据从工作进程 """
        return await aio.to_thread(self.output.get, timeout=3)
    
    async def register_process(self, process: mp.Process):
        """ 注册工作进程 """
        self.process = process
        if self.send_data_cache:
            for data in self.send_data_cache:
                await self.send(data)
            self.send_data_cache.clear()
        
    async def close(self):
        """ 关闭工作进程状态 """
        self._used = False
        await aio.to_thread(self.input.close)
        await aio.to_thread(self.output.close)
        if self.process is not None:
            await aio.to_thread(self.process.terminate)
            await aio.to_thread(self.process.join)
            logger.runner.info(f"工作进程 {self.worker_id} 已退出")
            self.process = None


class Context(BaseContext):
    """ 上下文 """
    @staticmethod
    async def try_(handler: Coroutine) -> None:
        """ 尝试执行事件处理函数 """
        try:
            await handler
        except Exception as e:
            logger.runner.exception(f"处理事件失败: {e}")
            
    @staticmethod
    def run_worker(worker_cls: Type[Worker], worker_name: str, settings: Settings, input: mp.Queue, output: mp.Queue):
        """ 运行工作进程 """
        setproctitle(worker_name)
        # 运行工作进程
        worker = worker_cls(worker_name)

        def exit_handler(_signum = None, _frame = None) -> None:
            """ 退出处理器 """
            logger.runner.info(f"工作进程 {worker_name} 收到退出信号，正在退出")
            # 清理日志
            logger.cleanup_loggers()
            worker.running = False
        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)
        atexit.register(exit_handler)

        worker.run(settings, input, output)
    
    def __init__(self, settings: Settings):
        self.running = True
        self.settings = settings
        self.workers: dict[str, WorkerState] = {}
        self.event_queue: aio.Queue[tuple[str, str, Any]] = aio.Queue()
        self.event_handlers: dict[str, list[Callable[[str, Any], Awaitable[None]]]] = {}
        self.tasks: TaskManager = TaskManager()
        
    def register_event_handler(self, event: str, handler: Callable[[str, Any], Awaitable[None]]):
        """ 注册事件处理函数 """
        if event not in self.event_handlers:
            self.event_handlers[event] = []
        self.event_handlers[event].append(handler)
        
    async def submit_event(self, source: str, event: str, data: Any):
        """ 提交事件 """
        self.event_queue.put_nowait((source, event, data))
        
    async def register_worker(self, worker_name: str, worker_id: int, worker_cls: Type[Worker]) -> WorkerState:
        """ 注册工作进程 """
        real_worker_id = f"lude-node:{worker_name}-{worker_id}"
        worker_state = WorkerState(real_worker_id, worker_cls)
        self.workers[real_worker_id] = worker_state
        await self.spawn_worker(worker_state)
        return worker_state
        
    async def spawn_worker(self, state: WorkerState):
        """ 启动工作进程 """
        process = mp.Process(
            target=self.run_worker, 
            args=(state.worker_cls, state.worker_id, self.settings, state.input, state.output),
            name=state.worker_id,
            daemon=True,
        )
        await state.register_process(process)
        process.start()
        await self.submit_event('Main', Events.WORKER_START, { 
            "worker_id": state.worker_id, 
            "process_id": process.pid,
        })
        logger.runner.info(f"启动工作进程: {state.worker_id}")
        
    async def keep_alive(self):
        """ 保持进程存活 """
        while self.running:
            # 检查 worker 进程状态
            for state in list(self.workers.values()):
                if not state.is_alive():
                    if state.process:
                        logger.runner.warning(f"工作进程 {state.worker_id} 未运行，准备重启")
                        await self.submit_event('Main', Events.WORKER_EXIT, {
                            "worker_id": state.worker_id,
                            "process_id": state.process.pid,
                        })
                    # mp.Queue 在子进程异常崩溃后无法安全复用（内部管道/锁可能损坏），
                    # 重启前必须关闭旧 Queue 并重建，否则新进程的 get()/主进程的 put()
                    # 会静默失败，导致进程虽在但完全不处理任务。
                    try:
                        state.input.close()
                        state.output.close()
                    except Exception:
                        pass
                    state.input = mp.Queue()
                    state.output = mp.Queue()
                    # 重新创建一个 worker
                    await self.spawn_worker(state)
            await aio.sleep(1)
        # 关闭所有工作进程
        if self.workers:
            await aio.wait([aio.create_task(state.close()) for state in self.workers.values()])
        self.workers.clear()
            
    async def event_loop(self):
        """ 事件循环 """
        while self.running:
            try:
                if event_data := await aio.wait_for(self.event_queue.get(), timeout=3):
                    source, event, data = event_data
                    # logger.runner.info(f"收到事件: {event}, 来源: {source}")
                    if event == Events.SHUTDOWN:
                        logger.runner.info("收到关闭事件, 正在退出事件循环")
                        self.running = False
                        break
                    if event in self.event_handlers:
                        for handle in self.event_handlers[event]:
                            self.tasks.create_task(self.try_(handle(source, data)))
            except aio.TimeoutError:
                continue
            
    async def notify_shutdown(self):
        """ 通知关闭上下文 """
        await self.event_queue.put(("Main", Events.SHUTDOWN, None))
