import logging, asyncio, json
from typing import Any, Callable, Awaitable, Self

from ._core import JsonRpcConnection
from ._exceptions import *

logger = logging.getLogger(__name__)

class JsonRpcStreamConnection(JsonRpcConnection):
    def __init__(self, host: str, port: int, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        super().__init__()
        self._host = host
        self._port = port
        self._reader = reader
        self._writer = writer

    @classmethod
    async def connect(cls, host: str, port: int, **kwargs) -> Self:
        """
        Client-side: connects to a JSON-RPC server over TCP.

        Args:
            host: The server host.
            port: The server port.
        """
        logger.info("Connecting to JSON-RPC TCP server at %s:%s", host, port)
        (reader, writer) = await asyncio.open_connection(host=host, port=port)
        connection = cls(host, port, reader, writer, **kwargs)
        connection._start_read_loop()
        return connection

    @classmethod
    async def serve(cls, host: str, port: int, *, handler: Callable[[Self], Awaitable] | None = None, **kwargs):
        """
        Server-side: starts a JSON-RPC TCP server.

        Args:
            host: The host to bind the server to.
            port: The port to bind the server to.
            handler: Async function called with each new connection.
        """
        if handler is None:
            async def default_handler(conn: Self):
                await conn.wait_closed()
            handler = default_handler
        logger.info("Starting JSON-RPC TCP server on %s:%s", host, port)
        async def on_connect(reader, writer):
            logger.debug("New connection from %s", writer.get_extra_info('peername'))
            conn = cls(host, port, reader, writer, **kwargs)
            conn._start_read_loop()
            await handler(conn)

        server = await asyncio.start_server(on_connect, host, port)
        async with server:
            await server.serve_forever()

    @property
    def host(self) -> str:
        return self._host

    @property
    def port(self) -> int:
        return self._port
    
    async def _send_bytes(self, data: bytes) -> None:
        self._writer.write(data + b"\n")
        await self._writer.drain()
            
    async def _cleanup(self) -> None:
        try:
            self._writer.close()
            await self._writer.wait_closed()
        except OSError:
            pass

    async def _read_loop(self) -> None:
        try:
            while True:
                line = await self._reader.readline()
                if line.strip() == b"":
                    logger.debug("TCP connection closed")
                    break
                obj = json.loads(line)
                await self._handle_incoming_message(obj)
        except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError, asyncio.IncompleteReadError):
            logger.debug("TCP connection lost")
        except Exception:
            logger.error("Error in read loop", exc_info=True)
        finally:
            await self.close()

__all__ = ["JsonRpcStreamConnection"]
