"""Cross platform abstractions for inter-process communication

On Unix, this uses AF_UNIX sockets.
On Windows, this uses NamedPipes.
"""

from __future__ import annotations

import base64
import codecs
import json
import os
import shutil
import sys
import tempfile
from abc import abstractmethod
from collections.abc import Callable
from select import select
from types import TracebackType
from typing import Final
from typing_extensions import Self

from librt.internal import ReadBuffer, WriteBuffer

if sys.platform == "win32":
    # This may be private, but it is needed for IPC on Windows, and is basically stable
    import _winapi
    import ctypes

    _IPCHandle = int

    kernel32 = ctypes.windll.kernel32
    DisconnectNamedPipe: Callable[[_IPCHandle], int] = kernel32.DisconnectNamedPipe
    FlushFileBuffers: Callable[[_IPCHandle], int] = kernel32.FlushFileBuffers
else:
    import socket

    _IPCHandle = socket.socket


class IPCException(Exception):
    """Exception for IPC issues."""


class IPCBase:
    """Base class for communication between the dmypy client and server.

    This contains logic shared between the client and server, such as reading
    and writing.
    We want to be able to send multiple "messages" over a single connection and
    to be able to separate the messages. We do this by encoding the messages
    in an alphabet that does not contain spaces, then adding a space for
    separation. The last framed message is also followed by a space.
    """

    connection: _IPCHandle

    def __init__(self, name: str, timeout: float | None) -> None:
        self.name = name
        self.timeout = timeout
        self.buffer = bytearray()

    def frame_from_buffer(self) -> bytearray | None:
        """Return a full frame from the bytes we have in the buffer."""
        space_pos = self.buffer.find(b" ")
        if space_pos == -1:
            return None
        # We have a full frame
        bdata = self.buffer[:space_pos]
        self.buffer = self.buffer[space_pos + 1 :]
        return bdata

    def read(self, size: int = 100000) -> str:
        return self.read_bytes(size).decode("utf-8")

    def read_bytes(self, size: int = 100000) -> bytes:
        """Read bytes from an IPC connection until we have a full frame."""
        bdata: bytearray | None = bytearray()
        if sys.platform == "win32":
            while True:
                # Check if we already have a message in the buffer before
                # receiving any more data from the socket.
                bdata = self.frame_from_buffer()
                if bdata is not None:
                    break

                # Receive more data into the buffer.
                ov, err = _winapi.ReadFile(self.connection, size, overlapped=True)
                try:
                    if err == _winapi.ERROR_IO_PENDING:
                        timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
                        res = _winapi.WaitForSingleObject(ov.event, timeout)
                        if res != _winapi.WAIT_OBJECT_0:
                            raise IPCException(f"Bad result from I/O wait: {res}")
                except BaseException:
                    ov.cancel()
                    raise
                _, err = ov.GetOverlappedResult(True)
                more = ov.getbuffer()
                if more:
                    self.buffer.extend(more)
                    bdata = self.frame_from_buffer()
                    if bdata is not None:
                        break
                if err == 0:
                    # we are done!
                    break
                elif err == _winapi.ERROR_MORE_DATA:
                    # read again
                    continue
                elif err == _winapi.ERROR_OPERATION_ABORTED:
                    raise IPCException("ReadFile operation aborted.")
        else:
            while True:
                # Check if we already have a message in the buffer before
                # receiving any more data from the socket.
                bdata = self.frame_from_buffer()
                if bdata is not None:
                    break

                # Receive more data into the buffer.
                more = self.connection.recv(size)
                if not more:
                    # Connection closed
                    break
                self.buffer.extend(more)

        if not bdata:
            # Socket was empty and we didn't get any frame.
            # This should only happen if the socket was closed.
            return b""
        return codecs.decode(bdata, "base64")

    def write(self, data: str) -> None:
        self.write_bytes(data.encode("utf-8"))

    def write_bytes(self, data: bytes) -> None:
        """Write to an IPC connection."""

        # Frame the data by urlencoding it and separating by space.
        encoded_data = codecs.encode(data, "base64") + b" "

        if sys.platform == "win32":
            try:
                ov, err = _winapi.WriteFile(self.connection, encoded_data, overlapped=True)
                try:
                    if err == _winapi.ERROR_IO_PENDING:
                        timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
                        res = _winapi.WaitForSingleObject(ov.event, timeout)
                        if res != _winapi.WAIT_OBJECT_0:
                            raise IPCException(f"Bad result from I/O wait: {res}")
                    elif err != 0:
                        raise IPCException(f"Failed writing to pipe with error: {err}")
                except BaseException:
                    ov.cancel()
                    raise
                bytes_written, err = ov.GetOverlappedResult(True)
                assert err == 0, err
                assert bytes_written == len(encoded_data)
            except OSError as e:
                raise IPCException(f"Failed to write with error: {e.winerror}") from e
        else:
            self.connection.sendall(encoded_data)

    def close(self) -> None:
        if sys.platform == "win32":
            if self.connection != _winapi.NULL:
                _winapi.CloseHandle(self.connection)
        else:
            self.connection.close()


class IPCClient(IPCBase):
    """The client side of an IPC connection."""

    def __init__(self, name: str, timeout: float | None) -> None:
        super().__init__(name, timeout)
        if sys.platform == "win32":
            timeout = int(self.timeout * 1000) if self.timeout else _winapi.NMPWAIT_WAIT_FOREVER
            try:
                _winapi.WaitNamedPipe(self.name, timeout)
            except FileNotFoundError as e:
                raise IPCException(f"The NamedPipe at {self.name} was not found.") from e
            except OSError as e:
                if e.winerror == _winapi.ERROR_SEM_TIMEOUT:
                    raise IPCException("Timed out waiting for connection.") from e
                else:
                    raise
            try:
                self.connection = _winapi.CreateFile(
                    self.name,
                    _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
                    0,
                    _winapi.NULL,
                    _winapi.OPEN_EXISTING,
                    _winapi.FILE_FLAG_OVERLAPPED,
                    _winapi.NULL,
                )
            except OSError as e:
                if e.winerror == _winapi.ERROR_PIPE_BUSY:
                    raise IPCException("The connection is busy.") from e
                else:
                    raise
            _winapi.SetNamedPipeHandleState(
                self.connection, _winapi.PIPE_READMODE_MESSAGE, None, None
            )
        else:
            self.connection = socket.socket(socket.AF_UNIX)
            self.connection.settimeout(timeout)
            self.connection.connect(name)

    def __enter__(self) -> IPCClient:
        return self

    def __exit__(
        self,
        exc_ty: type[BaseException] | None = None,
        exc_val: BaseException | None = None,
        exc_tb: TracebackType | None = None,
    ) -> None:
        self.close()


class IPCServer(IPCBase):
    BUFFER_SIZE: Final = 2**16

    def __init__(self, name: str, timeout: float | None = None) -> None:
        if sys.platform == "win32":
            name = r"\\.\pipe\{}-{}.pipe".format(
                name, base64.urlsafe_b64encode(os.urandom(6)).decode()
            )
        else:
            name = f"{name}.sock"
        super().__init__(name, timeout)
        if sys.platform == "win32":
            self.connection = _winapi.CreateNamedPipe(
                self.name,
                _winapi.PIPE_ACCESS_DUPLEX
                | _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
                | _winapi.FILE_FLAG_OVERLAPPED,
                _winapi.PIPE_READMODE_MESSAGE
                | _winapi.PIPE_TYPE_MESSAGE
                | _winapi.PIPE_WAIT
                | 0x8,  # PIPE_REJECT_REMOTE_CLIENTS
                1,  # one instance
                self.BUFFER_SIZE,
                self.BUFFER_SIZE,
                _winapi.NMPWAIT_WAIT_FOREVER,
                0,  # Use default security descriptor
            )
            if self.connection == -1:  # INVALID_HANDLE_VALUE
                err = _winapi.GetLastError()
                raise IPCException(f"Invalid handle to pipe: {err}")
        else:
            self.sock_directory = tempfile.mkdtemp()
            sockfile = os.path.join(self.sock_directory, self.name)
            self.sock = socket.socket(socket.AF_UNIX)
            self.sock.bind(sockfile)
            self.sock.listen(1)
            if timeout is not None:
                self.sock.settimeout(timeout)

    def __enter__(self) -> IPCServer:
        if sys.platform == "win32":
            # NOTE: It is theoretically possible that this will hang forever if the
            # client never connects, though this can be "solved" by killing the server
            try:
                ov = _winapi.ConnectNamedPipe(self.connection, overlapped=True)
            except OSError as e:
                # Don't raise if the client already exists, or the client already connected
                if e.winerror not in (_winapi.ERROR_PIPE_CONNECTED, _winapi.ERROR_NO_DATA):
                    raise
            else:
                try:
                    timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
                    res = _winapi.WaitForSingleObject(ov.event, timeout)
                    assert res == _winapi.WAIT_OBJECT_0
                except BaseException:
                    ov.cancel()
                    _winapi.CloseHandle(self.connection)
                    raise
                _, err = ov.GetOverlappedResult(True)
                assert err == 0
        else:
            try:
                self.connection, _ = self.sock.accept()
            except TimeoutError as e:
                raise IPCException("The socket timed out") from e
        return self

    def __exit__(
        self,
        exc_ty: type[BaseException] | None = None,
        exc_val: BaseException | None = None,
        exc_tb: TracebackType | None = None,
    ) -> None:
        if sys.platform == "win32":
            try:
                # Wait for the client to finish reading the last write before disconnecting
                if not FlushFileBuffers(self.connection):
                    raise IPCException(
                        "Failed to flush NamedPipe buffer, maybe the client hung up?"
                    )
            finally:
                DisconnectNamedPipe(self.connection)
        else:
            self.close()

    def cleanup(self) -> None:
        if sys.platform == "win32":
            self.close()
        else:
            shutil.rmtree(self.sock_directory)

    @property
    def connection_name(self) -> str:
        if sys.platform == "win32":
            return self.name
        elif sys.platform == "gnu0":
            # GNU/Hurd returns empty string from getsockname()
            # for AF_UNIX sockets
            return os.path.join(self.sock_directory, self.name)
        else:
            name = self.sock.getsockname()
            assert isinstance(name, str)
            return name


class BadStatus(Exception):
    """Exception raised when there is something wrong with the status file.

    For example:
    - No status file found
    - Status file malformed
    - Process whose pid is in the status file does not exist
    """


def read_status(status_file: str) -> dict[str, object]:
    """Read status file.

    Raise BadStatus if the status file doesn't exist or contains
    invalid JSON or the JSON is not a dict.
    """
    if not os.path.isfile(status_file):
        raise BadStatus("No status file found")
    with open(status_file) as f:
        try:
            data = json.load(f)
        except Exception as e:
            raise BadStatus("Malformed status file (not JSON)") from e
    if not isinstance(data, dict):
        raise BadStatus("Invalid status file (not a dict)")
    return data


def ready_to_read(conns: list[IPCClient], timeout: float | None = None) -> list[int]:
    """Wait until some connections are readable.

    Return index of each readable connection in the original list.
    """
    # TODO: add Windows support for this.
    assert sys.platform != "win32"
    connections = [conn.connection for conn in conns]
    ready, _, _ = select(connections, [], [], timeout)
    return [connections.index(r) for r in ready]


def send(connection: IPCBase, data: IPCMessage) -> None:
    """Send data to a connection encoded and framed.

    The data must be a non-abstract IPCMessage. We assume that a single send call is a
    single frame to be sent.
    """
    buf = WriteBuffer()
    data.write(buf)
    connection.write_bytes(buf.getvalue())


def receive(connection: IPCBase) -> ReadBuffer:
    """Receive single encoded IPCMessage frame from a connection.

    Raise OSError if the data received is not valid.
    """
    bdata = connection.read_bytes()
    if not bdata:
        raise OSError("No data received")
    return ReadBuffer(bdata)


class IPCMessage:
    @classmethod
    @abstractmethod
    def read(cls, buf: ReadBuffer) -> Self:
        raise NotImplementedError

    @abstractmethod
    def write(self, buf: WriteBuffer) -> None:
        raise NotImplementedError
