import itertools, logging, asyncio, json, inspect
from typing import Any, Callable, Self
from abc import ABC, abstractmethod
from ._exceptions import *

logger = logging.getLogger(__name__)

def jsonrpc_method(name: str | None = None):
    def decorator(fn):
        fn._is_rpc_method = True
        fn._rpc_method_name = name or fn.__name__
        return fn
    return decorator

class JsonRpcMethod:
    def __init__(self, func: Callable[..., Any], name: str | None = None, is_class_method: bool = False):
        self.func = func
        self.signature = inspect.signature(func)
        self.name = name or func.__name__
        self.is_class_method = is_class_method

    async def call(self, connection, params):
        if isinstance(params, list):
            args, kwargs = params, {}
        elif isinstance(params, dict):
            args, kwargs = [], params
        else:
            args, kwargs = [], {}

        if self.is_class_method:
            bound = self.signature.bind(connection, *args, **kwargs) # the args correctly have connection :)
        else:
            bound = self.signature.bind(*args, **kwargs) # the args are passed as-is, connection is not expected

        result = self.func(*bound.args, **bound.kwargs)
        if asyncio.iscoroutine(result):
            result = await result
        return result

class JsonRpcConnection(ABC):
    def __init__(self):
        self._closed = False
        self._close_event = asyncio.Event()
        self._id_counter = itertools.count()

        self._send_lock = asyncio.Lock()

        self._pending: dict[int | str | float, asyncio.Queue] = {}
        self._methods: dict[str, JsonRpcMethod] = {
            obj._rpc_method_name: JsonRpcMethod(obj, obj._rpc_method_name, is_class_method=True)
            for _, obj in inspect.getmembers(self.__class__)
            if getattr(obj, "_is_rpc_method", False)
        }

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        await self.close()

    @classmethod
    @abstractmethod
    async def connect(cls, *args, **kwargs) -> Self: ...

    @classmethod
    @abstractmethod
    async def serve(cls, *args, **kwargs): ...

    @abstractmethod
    async def _send_bytes(self, data: bytes) -> None: ...
   
    @abstractmethod
    async def _cleanup(self) -> None: ...

    @abstractmethod
    async def _read_loop(self) -> None: ...

    def _start_read_loop(self) -> None:
        self._read_task = asyncio.create_task(self._read_loop())
        
    async def wait_closed(self) -> None:
        """Waits until the connection is closed."""
        await self._close_event.wait()

    async def call(
        self, method: str, params: Any, timeout: float | None = 30.0
    ) -> Any:
        """Calls a JSON-RPC method and waits for the response.

        Args:
            method: The name of the method to call.
            params: Parameters to pass to the method.
            timeout: Optional timeout value (in seconds) before raising an exception.

        Returns:
            The result from the method call.

        Raises:
            JsonRpcException: If the remote method returns an error response.
            ServerErrorException: If the connection is closed, send fails, or the response is malformed.
        """
        rid = self._next_id()
        msg = {"jsonrpc": "2.0", "method": method, "id": rid}
        if params is not None:
            msg["params"] = params

        q: asyncio.Queue = asyncio.Queue(maxsize=1)
        if self._closed:
            raise ServerErrorException("Connection closed")
        self._pending[rid] = q

        try:
            await self._send_json(msg)
        except Exception:
            self._pending.pop(rid, None)
            raise ServerErrorException("Failed to send request")

        try:
            resp = await asyncio.wait_for(q.get(), timeout=timeout)
        except asyncio.TimeoutError:
            self._pending.pop(rid, None)
            raise ServerErrorException(f"Request {rid} timed out")

        if "error" in resp:
            logger.warning("Error response for id %s: %s", rid, resp["error"])
            err = resp["error"]
            raise JsonRpcException.from_error(err)
        if "result" in resp:
            return resp["result"]
        raise ServerErrorException("Invalid response")

    async def notify(
        self, method: str, params: Any
    ) -> None:
        """Sends a JSON-RPC notification (no response expected).

        Args:
            method: The name of the method to call.
            params: Parameters to pass to the method.

        Raises:
            ServerErrorException: If the connection is closed or send fails.
        """
        msg: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            msg["params"] = params
        try:
            await self._send_json(msg)
        except Exception:
            raise ServerErrorException("Failed to send notification")

    def add_method(self, func: Callable[..., Any], name: str | None = None):
        """Registers a method (at runtime) that can be called by the remote side.

        Args:
            name: Optional name for the method, by default the function name is used.
        """
        logger.debug("Registering method '%s'", name)
        real_name = name or func.__name__
        self._methods[real_name] = JsonRpcMethod(func, name, is_class_method=False)

    def remove_method(self, method: str) -> None:
        """Removes a previously registered method.

        Args:
            method: The name of the method to remove.
        """
        self._methods.pop(method, None)

    async def close(self) -> None:
        """Closes the connection and fails all pending calls."""
        if self._closed:
            return
        self._closed = True
        logger.debug("Closing connection...")
        if self._pending:
            logger.warning("Failing %d pending calls", len(self._pending))

        # cancelling the read loop safely...
        if asyncio.current_task() is not self._read_task:
            self._read_task.cancel()
            try:
                await self._read_task
            except (asyncio.CancelledError, Exception):
                pass        

        for q in self._pending.values():
            q.put_nowait({"error": {"code": -32000, "message": "Connection closed"}})
        self._pending.clear()
        self._close_event.set()
        await self._cleanup()

    @property
    def is_closed(self) -> bool:
        """Returns True if the connection is closed / dead."""
        return self._closed

    def _next_id(self) -> int:
        return next(self._id_counter)

    async def _send_json(self, message: dict[str, Any]) -> None:
        data = json.dumps(message).encode()
        logger.debug("sending JSON-RPC object")
        async with self._send_lock:
            try:
                await self._send_bytes(data)
            except Exception:
                await self.close()
                raise

    async def _send_error(self, error: JsonRpcException, rid) -> None:
        error_dict = error.to_json_rpc_error()
        await self._send_json(
            {"jsonrpc": "2.0", "id": rid, "error": error_dict}
        )

    async def _deliver_response(self, rid: int | str | float, obj: dict[str, Any]) -> None:
        q = self._pending.pop(rid, None)
        if q:
            await q.put(obj)
            logger.debug("Delivered response for id %s: %s", rid, obj)
        else:
            logger.warning("Received response for unknown id %s: %s", rid, obj)

    async def _handle_incoming_message(self, obj) -> None:
        logger.debug("Received JSON-RPC object: %s", obj)
        if isinstance(obj, list):
            incoming = obj
        elif isinstance(obj, dict):
            incoming = [obj]
        else:
            await self._send_error(InvalidRequestException(), None)
            return

        for obj in incoming:
            rid = obj.get("id")
            if isinstance(rid, bool):
                await self._send_error(InvalidRequestException(), None)
                continue
            if rid is not None and not isinstance(rid, (str, int, float)):
                await self._send_error(InvalidRequestException(), None)
                continue 

            if obj.get("jsonrpc") != "2.0":
                await self._send_error(InvalidRequestException(), rid)
                continue 

            if rid is not None and "method" not in obj:
                await self._deliver_response(rid, obj)
                continue 

            if "method" in obj:
                method_name = obj["method"]
                params = obj.get("params")
                method = self._methods.get(method_name)

                if not method:
                    if rid is not None:
                        await self._send_error(MethodNotFoundException(method_name), rid)
                    continue 
                try:
                    result = await method.call(self, params)

                    if rid is not None:
                        await self._send_json({"jsonrpc": "2.0", "id": rid, "result": result})
                except TypeError as e:
                    if rid is not None:
                        await self._send_error(InvalidParamsException(str(e)), rid)
                except JsonRpcException as e:
                    if rid is not None:
                        await self._send_error(e, rid)
                except Exception as e:
                    logger.error("Error calling method '%s'", method, exc_info=True)
                    if rid is not None:
                        await self._send_error(InternalErrorException(str(e)), rid)
                continue

            await self._send_error(InvalidRequestException(), obj.get("id"))

__all__ = ["JsonRpcConnection", "jsonrpc_method"]
