from __future__ import annotations

import asyncio
from dataclasses import dataclass
from enum import Enum, auto
from itertools import count
from typing import Any, Awaitable, Callable, Optional


class OpState(Enum):
    PENDING = auto()
    RUNNING = auto()
    DONE = auto()
    FAILED = auto()
    CANCELLED = auto()


@dataclass
class OpInfo:
    state: OpState
    future: asyncio.Future
    result: Any = None
    error: Optional[BaseException] = None


class AsyncRW:
    """
    - Two queues: read_q, write_q
    - Two worker tasks created at start()
    - read()/write() enqueue and return op_id
    - poll status via status(op_id) / done(op_id) / result(op_id)
    """

    def __init__(
        self,
        read_handler: Callable[[Any], Awaitable[Any]],
        write_handler: Callable[[Any], Awaitable[Any]],
        *,
        max_queue: int = 0,
    ) -> None:
        self._read_handler = read_handler
        self._write_handler = write_handler

        self._read_q: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(maxsize=max_queue)
        self._write_q: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(maxsize=max_queue)

        self._ops: dict[str, OpInfo] = {}
        self._id_read = count(1)
        self._id_write = count(1)

        self._stop = asyncio.Event()
        self._read_worker: Optional[asyncio.Task] = None
        self._write_worker: Optional[asyncio.Task] = None

    # --- lifecycle ---------------------------------------------------------

    async def start(self) -> None:
        if self._read_worker or self._write_worker:
            return
        self._stop.clear()
        self._read_worker = asyncio.create_task(self._run_read_worker(), name="rw:read-worker")
        self._write_worker = asyncio.create_task(self._run_write_worker(), name="rw:write-worker")

    async def close(self) -> None:
        self._stop.set()
        # Cancel workers so they exit promptly even if blocked on .get()
        for t in (self._read_worker, self._write_worker):
            if t:
                t.cancel()
        await asyncio.gather(
            *(t for t in (self._read_worker, self._write_worker) if t),
            return_exceptions=True,
        )
        self._read_worker = None
        self._write_worker = None

    @property
    def read_worker_id(self) -> Optional[str]:
        return self._read_worker.get_name() if self._read_worker else None

    @property
    def write_worker_id(self) -> Optional[str]:
        return self._write_worker.get_name() if self._write_worker else None

    # --- enqueue ops -------------------------------------------------------

    def read(self, req: Any) -> str:
        op_id = f"r{next(self._id_read)}"
        fut = asyncio.get_running_loop().create_future()
        self._ops[op_id] = OpInfo(state=OpState.PENDING, future=fut)
        self._read_q.put_nowait((op_id, req))
        return op_id

    def write(self, item: Any) -> str:
        op_id = f"w{next(self._id_write)}"
        fut = asyncio.get_running_loop().create_future()
        self._ops[op_id] = OpInfo(state=OpState.PENDING, future=fut)
        self._write_q.put_nowait((op_id, item))
        return op_id

    # --- polling / results -------------------------------------------------

    def status(self, op_id: str) -> OpState:
        return self._ops[op_id].state

    def done(self, op_id: str) -> bool:
        return self._ops[op_id].future.done()

    async def wait(self, op_id: str) -> Any:
        return await self._ops[op_id].future

    def result(self, op_id: str) -> Any:
        # raises if not done / failed, like Future.result()
        return self._ops[op_id].future.result()

    def exception(self, op_id: str) -> Optional[BaseException]:
        return self._ops[op_id].future.exception()

    def cancel(self, op_id: str) -> bool:
        info = self._ops.get(op_id)
        if not info:
            return False
        info.state = OpState.CANCELLED
        return info.future.cancel()

    # --- workers -----------------------------------------------------------

    async def _run_read_worker(self) -> None:
        try:
            while not self._stop.is_set():
                op_id, req = await self._read_q.get()
                await self._execute(op_id, self._read_handler, req)
                self._read_q.task_done()
        except asyncio.CancelledError:
            pass

    async def _run_write_worker(self) -> None:
        try:
            while not self._stop.is_set():
                op_id, item = await self._write_q.get()
                await self._execute(op_id, self._write_handler, item)
                self._write_q.task_done()
        except asyncio.CancelledError:
            pass

    async def _execute(
        self,
        op_id: str,
        handler: Callable[[Any], Awaitable[Any]],
        payload: Any,
    ) -> None:
        info = self._ops.get(op_id)
        if not info or info.future.cancelled():
            return

        info.state = OpState.RUNNING
        try:
            res = await handler(payload)
            info.result = res
            info.state = OpState.DONE
            if not info.future.done():
                info.future.set_result(res)
        except asyncio.CancelledError:
            info.state = OpState.CANCELLED
            if not info.future.done():
                info.future.cancel()
        except BaseException as e:
            info.error = e
            info.state = OpState.FAILED
            if not info.future.done():
                info.future.set_exception(e)