from dataclasses import dataclass
from logging import getLogger
from multiprocessing import Pool
from multiprocessing.pool import Pool as PoolClass
from socket import socket
from socketserver import BaseRequestHandler, StreamRequestHandler, ThreadingUnixStreamServer, ThreadingTCPServer
from typing import Any, Callable, Self
import json
import os
import struct

from .handle_dict import render_from_dict

LOGGER = getLogger(__name__)


@dataclass
class ServiceArguments:
    """Arguments to the service"""

    pool_size: int
    unix_socket_path: str
    socket_timeout: int | None


class NoDataError(Exception):
    """Error for no data on the socket - app should not write an error to the socket"""

    def __init__(self) -> None:
        super().__init__("No data on socket")


class InsufficientDataError(Exception):
    """Error for insufficient data on the socket - app should not write an error"""

    def __init__(self, expected: int, received: bytes, hint: bool) -> None:
        self.expected = expected
        self.received = received
        super().__init__(f"expected {expected} bytes, got {len(received)}" + f" '{received.hex()}'" if hint else "")


class RequestHandler(StreamRequestHandler):
    """A request handler for the renderer protocol"""

    # one pool shared across all users
    pool: PoolClass | None = None

    @classmethod
    def setup_pool(cls, size: int) -> None:
        """Set up the pool shared across all the handlers"""
        if cls.pool is None:
            cls.pool = Pool(size)

    @classmethod
    def teardown_pool(cls) -> None:
        """Tear down the pool shared across all the handlers"""
        if cls.pool is not None:
            cls.pool.terminate()
            cls.pool.join()
            cls.pool = None

    def write_message(self, data):
        """JSON encode and write a message to the socket"""
        result_as_bytes = json.dumps(data).encode("utf-8")
        size = struct.pack(">L", len(result_as_bytes))
        response = size + result_as_bytes
        self.wfile.write(response)

    def write_error(self, error: str):
        """Write an error message to the socket"""
        try:
            self.write_message({"error": {"errorString": error, "errorDetail": ""}})
        except Exception as exc:
            LOGGER.error(f"Got an error writing error message: {exc}")  # oh well

    def read_bytes(self, num_bytes: int) -> bytes:
        """Read num_bytes bytes off the socket, possibly performing multiple reads"""
        data_seen = b""
        while len(data_seen) < num_bytes:
            last_read = self.rfile.read(num_bytes - len(data_seen))
            if not last_read:
                # only show the hint if we expected a small # of bytes
                raise InsufficientDataError(num_bytes, data_seen, hint=len(data_seen) < 64)
            data_seen += last_read
        return data_seen

    def read_message(self):
        """Read one JSON message off the socket"""
        LOGGER.debug("reading 4 bytes")
        size_as_bytes = self.rfile.read(4)
        if not size_as_bytes:
            raise NoDataError()
        elif len(size_as_bytes) != 4:
            # we should always get an atomic 0 or 4 bytes on a unix domain socket for this read, but...
            raise InsufficientDataError(4, size_as_bytes, True)
        (size,) = struct.unpack(">L", size_as_bytes)
        LOGGER.debug(f"got size (bytes: {size_as_bytes.hex()}), reading {size} bytes")
        data = self.read_bytes(size)
        LOGGER.debug(f"read {len(data)} bytes")
        request_data = json.loads(data.decode("utf-8"))
        return request_data

    def handle(self) -> None:
        """Handle a single message on the socket"""
        if not self.pool:
            raise Exception("Pool not initialized!")
        try:
            request_data = self.read_message()
            result = self.pool.apply(render_from_dict, [request_data])
        except NoDataError as exc:
            LOGGER.error(exc)
            return
        except InsufficientDataError as exc:
            LOGGER.error(exc)
            self.write_error(str(exc))
            return
        except Exception as exc:
            LOGGER.error(f"Unexpected error: {exc}")
            self.write_error(str(exc))
            return
        # this needs to go into its own block: if it fails we don't want to try writing to the socket!
        try:
            self.write_message(result)
        except Exception as exc:
            LOGGER.error(f"Could not write response: {exc}")
            return


class TimeoutUnixStreamingServer(ThreadingUnixStreamServer):
    """A stream server that adds a timeout to sockets that come back from accept()."""

    def __init__(
        self,
        server_address: str,
        RequestHandlerClass: Callable[[Any, Any, Self], BaseRequestHandler],
        socket_timeout: int,
        bind_and_activate: bool = True,
    ) -> None:
        super().__init__(server_address, RequestHandlerClass, bind_and_activate)
        self.socket_timeout = socket_timeout

    def get_request(self) -> tuple[socket, Any]:
        (sock, addr) = super().get_request()
        sock.settimeout(self.socket_timeout)  # set a 30s timeout on the socket
        return sock, addr


def serve(args: ServiceArguments):
    """The core serve loop"""
    RequestHandler.setup_pool(args.pool_size)
    server: TimeoutUnixStreamingServer | ThreadingTCPServer
    LOGGER.info(f"starting with path={args.unix_socket_path}")
    server = TimeoutUnixStreamingServer(args.unix_socket_path, RequestHandler, socket_timeout=args.socket_timeout or 30)
    try:
        server.serve_forever()
    finally:
        RequestHandler.teardown_pool()
        if args.unix_socket_path:
            os.remove(args.unix_socket_path)
