"""Spawned model processes, reused across tasks and replaceable after a crash."""
import logging
import multiprocessing
import queue
import time

logger = logging.getLogger(__name__)


def analysis_process(connection, model_path, workers):
    try:
        from wanzheng import SafeYOLO, analyze_directory
        model = SafeYOLO(model_path)
        connection.send({'ready': True})
        while True:
            task = connection.recv()
            if task is None:
                return
            result = analyze_directory(task['input_dir'], task['task_id'], model,
                                       workers, task.get('api_url'))
            connection.send(result)
    except EOFError:
        pass
    except BaseException as exc:
        logger.exception('常驻分析进程异常')
        try:
            connection.send({'error': str(exc) or type(exc).__name__})
        except (OSError, EOFError):
            pass
    finally:
        connection.close()


class ResidentPool:
    def __init__(self, model_path, concurrency=1, workers=4, api_url=None,
                 startup_timeout=180, task_timeout=3600, target=analysis_process):
        if min(concurrency, workers, startup_timeout, task_timeout) <= 0:
            raise ValueError('并发、线程数和超时必须大于 0')
        self.model_path = str(model_path)
        self.concurrency = concurrency
        self.workers = workers
        self.api_url = api_url
        self.startup_timeout = startup_timeout
        self.task_timeout = task_timeout
        self.target = target
        self.context = multiprocessing.get_context('spawn')
        self.slots = queue.Queue()

    def receive(self, slot, timeout, stop_event=None):
        connection, process = slot
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if stop_event is not None and stop_event.is_set():
                raise RuntimeError('服务停止，任务中断，请重新提交')
            if connection.poll(0.1):
                try:
                    return connection.recv()
                except (EOFError, OSError) as exc:
                    raise RuntimeError('常驻分析进程连接中断') from exc
            if not process.is_alive():
                raise RuntimeError(f'常驻分析进程退出：{process.exitcode}')
        raise RuntimeError('常驻分析进程等待超时')

    @staticmethod
    def dispose(slot):
        if slot is None:
            return
        connection, process = slot
        connection.close()
        if process.is_alive():
            process.terminate()
        process.join(timeout=5)
        if process.is_alive():
            process.kill()
            process.join()
        process.close()

    def spawn(self, stop_event=None):
        parent, child = self.context.Pipe()
        process = self.context.Process(target=self.target,
                                       args=(child, self.model_path, self.workers))
        try:
            process.start()
        except BaseException:
            parent.close()
            child.close()
            process.close()
            raise
        child.close()
        slot = parent, process
        try:
            message = self.receive(slot, self.startup_timeout, stop_event)
            if not message.get('ready'):
                raise RuntimeError(message.get('error', '模型预热失败'))
            return slot
        except BaseException:
            self.dispose(slot)
            raise

    def start(self):
        try:
            for _ in range(self.concurrency):
                self.slots.put(self.spawn())
        except BaseException:
            self.stop()
            raise

    def stop(self):
        # TaskQueue joins its worker threads before closing idle slots.
        while True:
            try:
                self.dispose(self.slots.get_nowait())
            except queue.Empty:
                break

    def __call__(self, task, stop_event):
        slot = self.slots.get()
        try:
            if stop_event.is_set():
                raise RuntimeError('服务停止，任务中断，请重新提交')
            if slot is None or not slot[1].is_alive():
                self.dispose(slot)
                slot = None
                slot = self.spawn(stop_event)
            slot[0].send(dict(task, api_url=self.api_url))
            result = self.receive(slot, self.task_timeout, stop_event)
        except BaseException:
            self.dispose(slot)
            slot = None
            raise
        finally:
            self.slots.put(slot)
        if result.get('error'):
            raise RuntimeError(result['error'])
        return result