from __future__ import annotations

import json
from pathlib import Path

from PySide6.QtCore import QObject, Signal
from PySide6.QtNetwork import QLocalServer, QLocalSocket

from bedit_gui.simulation_models import CompiledModel

SERVER_NAME = "bedit-besim-handoff-v1"


def send_simulation_handoff(path: str | Path, compiled_model: CompiledModel, *, server_name: str = SERVER_NAME, timeout_ms: int = 500) -> bool:
    socket = QLocalSocket()
    socket.connectToServer(server_name)
    if not socket.waitForConnected(timeout_ms):
        return False
    payload = json.dumps({"path": str(Path(path).resolve()), "compiled_model": compiled_model.to_data()}, separators=(",", ":")).encode("utf-8") + b"\n"
    if socket.write(payload) != len(payload) or not socket.waitForBytesWritten(timeout_ms):
        socket.abort()
        return False
    socket.disconnectFromServer()
    return True


class SimulationHandoffServer(QObject):
    handoff_received = Signal(object, object)

    def __init__(self, parent: QObject | None = None, *, server_name: str = SERVER_NAME) -> None:
        super().__init__(parent)
        self.server = QLocalServer(self)
        self._buffers: dict[QLocalSocket, bytearray] = {}
        self.server.newConnection.connect(self._accept_connections)
        if not self.server.listen(server_name):
            probe = QLocalSocket()
            probe.connectToServer(server_name)
            if probe.waitForConnected(200):
                probe.disconnectFromServer()
                return
            QLocalServer.removeServer(server_name)
            if not self.server.listen(server_name):
                raise RuntimeError(f"could not listen for BEsim handoffs: {self.server.errorString()}")

    def _accept_connections(self) -> None:
        while self.server.hasPendingConnections():
            socket = self.server.nextPendingConnection()
            if socket is None:
                continue
            self._buffers[socket] = bytearray()
            socket.readyRead.connect(lambda active=socket: self._read(active))
            socket.disconnected.connect(lambda active=socket: self._discard(active))
            self._read(socket)

    def _read(self, socket: QLocalSocket) -> None:
        buffer = self._buffers.get(socket)
        if buffer is None:
            return
        buffer.extend(socket.readAll().data())
        while b"\n" in buffer:
            raw_message, _, remaining = buffer.partition(b"\n")
            buffer[:] = remaining
            try:
                message = json.loads(raw_message.decode("utf-8"))
                path = Path(str(message["path"]))
                compiled_model = CompiledModel.from_data(message["compiled_model"])
            except (KeyError, TypeError, ValueError, json.JSONDecodeError, UnicodeDecodeError):
                continue
            self.handoff_received.emit(path, compiled_model)

    def _discard(self, socket: QLocalSocket) -> None:
        self._buffers.pop(socket, None)