"""Small, bounded, persistent subprocess worker for CPU/model plugins.

The worker is intentionally not wired automatically from YAML.  A model plugin
owns one worker for its whole lifecycle, loads its model once in the child, and
submits small payloads or shared-memory references.  Large decoded frames
should not be copied through these queues.
"""

from __future__ import annotations

import asyncio
import multiprocessing as mp
import queue
import traceback
from collections.abc import Callable
from typing import Any
from uuid import uuid4
from time import monotonic


WorkerHandler = Callable[[Any], Any]
WorkerFactory = Callable[[], WorkerHandler]
_STOP = "__cmvr_edge_ai_worker_stop__"


class ProcessWorkerError(RuntimeError):
    pass


class RemoteWorkerError(ProcessWorkerError):
    def __init__(self, remote_type: str, message: str, remote_traceback: str) -> None:
        self.remote_type = remote_type
        self.remote_traceback = remote_traceback
        super().__init__(f"remote {remote_type}: {message}\n{remote_traceback}")


class PersistentProcessWorker:
    """One long-lived subprocess with bounded request and result queues."""

    def __init__(
        self,
        factory: WorkerFactory,
        *,
        name: str = "model-worker",
        start_method: str = "spawn",
        queue_capacity: int = 1,
    ) -> None:
        if queue_capacity < 1:
            raise ValueError("queue_capacity must be at least 1")
        self._factory = factory
        self._name = name
        self._context: Any = mp.get_context(start_method)
        self._requests = self._context.Queue(maxsize=queue_capacity)
        self._results = self._context.Queue(maxsize=queue_capacity)
        self._process: mp.Process | None = None
        self._submit_lock = asyncio.Lock()
        self._closed = False
        self._poisoned_reason: str | None = None

    @property
    def is_alive(self) -> bool:
        return self._process is not None and self._process.is_alive()

    @property
    def poisoned(self) -> bool:
        """Whether an in-flight timeout/cancellation made correlation unsafe."""

        return self._poisoned_reason is not None

    async def start(self) -> None:
        if self._closed:
            raise ProcessWorkerError("worker is one-shot and has already been stopped")
        if self._process is not None:
            raise ProcessWorkerError("worker has already been started")
        process = self._context.Process(
            target=_worker_main,
            args=(self._factory, self._requests, self._results),
            name=self._name,
        )
        process.start()
        self._process = process

    async def submit(self, payload: Any, *, timeout_s: float | None = None) -> Any:
        if not self.is_alive:
            raise ProcessWorkerError("worker is not running")
        if self._poisoned_reason is not None:
            raise ProcessWorkerError(
                "worker is poisoned by an indeterminate prior request; stop and "
                f"recreate it ({self._poisoned_reason})"
            )
        if timeout_s is not None and timeout_s <= 0:
            raise ValueError("timeout_s must be positive")
        request_id = uuid4().hex
        async with self._submit_lock:
            if self._poisoned_reason is not None:
                raise ProcessWorkerError(
                    "worker is poisoned by an indeterminate prior request; stop "
                    f"and recreate it ({self._poisoned_reason})"
                )
            deadline = monotonic() + timeout_s if timeout_s is not None else None
            request_may_have_been_sent = False
            try:
                # Cancellation can race the backing queue thread, so once the
                # put begins we conservatively treat its delivery as uncertain.
                request_may_have_been_sent = True
                await self._put_request((request_id, payload), deadline)
                response = await self._get_response(deadline)
            except BaseException as exc:
                if request_may_have_been_sent:
                    self._poisoned_reason = type(exc).__name__
                raise

        response_id, success, value = response
        if response_id != request_id:
            self._poisoned_reason = "response correlation mismatch"
            raise ProcessWorkerError(
                f"worker response correlation mismatch: {response_id} != {request_id}"
            )
        if success:
            return value
        remote_type, message, remote_traceback = value
        raise RemoteWorkerError(remote_type, message, remote_traceback)

    async def _put_request(self, request: Any, deadline: float | None) -> None:
        while True:
            if not self.is_alive:
                raise ProcessWorkerError("worker exited before accepting request")
            timeout = _poll_timeout(deadline)
            try:
                await asyncio.to_thread(self._requests.put, request, True, timeout)
                return
            except queue.Full:
                if deadline is not None and monotonic() >= deadline:
                    raise TimeoutError("worker request timed out")

    async def _get_response(self, deadline: float | None) -> Any:
        while True:
            if not self.is_alive:
                raise ProcessWorkerError("worker exited before returning a response")
            timeout = _poll_timeout(deadline)
            try:
                return await asyncio.to_thread(self._results.get, True, timeout)
            except queue.Empty:
                if deadline is not None and monotonic() >= deadline:
                    raise TimeoutError("worker response timed out")

    async def stop(self, *, timeout_s: float = 5.0) -> None:
        if timeout_s <= 0:
            raise ValueError("timeout_s must be positive")
        process = self._process
        if process is None:
            return
        if process.is_alive():
            try:
                await asyncio.to_thread(
                    self._requests.put, _STOP, True, min(timeout_s, 1.0)
                )
            except queue.Full:
                pass
            await asyncio.to_thread(process.join, timeout_s)
            if process.is_alive():
                process.terminate()
                await asyncio.to_thread(process.join, 1.0)
        self._process = None
        self._requests.close()
        self._results.close()
        self._closed = True


def _worker_main(factory: WorkerFactory, requests: Any, results: Any) -> None:
    try:
        handler = factory()
    except BaseException as exc:
        # There is no request id yet, so the parent observes process death.
        traceback.print_exception(exc)
        return

    while True:
        item = requests.get()
        if item == _STOP:
            return
        request_id, payload = item
        try:
            result = handler(payload)
            results.put((request_id, True, result))
        except BaseException as exc:
            results.put(
                (
                    request_id,
                    False,
                    (type(exc).__name__, str(exc), traceback.format_exc()),
                )
            )


def _poll_timeout(deadline: float | None) -> float:
    if deadline is None:
        return 0.2
    remaining = deadline - monotonic()
    if remaining <= 0:
        return 0.0
    return min(remaining, 0.2)


__all__ = [
    "PersistentProcessWorker",
    "ProcessWorkerError",
    "RemoteWorkerError",
    "WorkerFactory",
    "WorkerHandler",
]
