"""IPC connection from the training process to the sidecar.

The connection is fork-safe: any write attempt from a forked child process
(e.g. a DataLoader worker) is a silent no-op, detected via PID comparison.
"""

import json
import os
import socket
import tempfile
from typing import Any


def _socket_path(*, client_uuid: str) -> str:
    """Return the Unix socket path for a given client UUID.

    Path is at most ~56 chars (well within the 104-byte macOS limit).
    Uses /tmp so the sidecar process (which may run as the same user)
    can always access it.
    """
    tmp = tempfile.gettempdir()
    return os.path.join(tmp, f"tyko-{client_uuid}.sock")


class IPCConnection:
    """Synchronous, fork-safe connection to the sidecar process.

    The connection is bound to the PID of the process that created it.
    Any call from a forked child process is a silent no-op.

    Example:
        conn = IPCConnection.connect(client_uuid="abc123")
        conn.send({"type": "log", "values": {...}, ...})      # fire-and-forget
        conn.send_and_wait({"type": "close", ...})             # waits for {"ok": true}
        conn.close()
    """

    def __init__(self, *, sock: socket.socket, client_uuid: str, owner_pid: int) -> None:
        self._sock = sock
        self._client_uuid = client_uuid
        self._owner_pid = owner_pid

    @classmethod
    def connect(cls, *, client_uuid: str, timeout: float = 2.0) -> "IPCConnection":
        """Connect to an existing sidecar socket.

        Raises:
            OSError: if the socket does not exist or connection is refused.
        """
        path = _socket_path(client_uuid=client_uuid)
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect(path)
        sock.settimeout(None)  # switch to blocking after connect
        return cls(sock=sock, client_uuid=client_uuid, owner_pid=os.getpid())

    def _is_owner(self) -> bool:
        return os.getpid() == self._owner_pid

    def send(self, message: dict[str, Any]) -> None:
        """Write a newline-delimited JSON message to the socket.

        Silent no-op if:
        - Called from a forked child process (PID mismatch)
        - The socket is broken (BrokenPipeError / ConnectionResetError)
        """
        if not self._is_owner():
            return
        try:
            data = json.dumps(message).encode() + b"\n"
            self._sock.sendall(data)
        except (BrokenPipeError, ConnectionResetError, OSError):
            pass

    def send_and_wait(self, message: dict[str, Any], *, timeout: float = 10.0) -> bool:
        """Send a message and block until {"ok": true} is received.

        Used for "set_run", "close", and "shutdown" messages to ensure
        the sidecar has processed them before the main process continues.

        Returns:
            True if the sidecar acknowledged, False if timeout exceeded
            or the socket is broken.
        """
        if not self._is_owner():
            return False
        try:
            data = json.dumps(message).encode() + b"\n"
            self._sock.sendall(data)
            self._sock.settimeout(timeout)
            buf = b""
            while b"\n" not in buf:
                chunk = self._sock.recv(4096)
                if not chunk:
                    return False
                buf += chunk
            self._sock.settimeout(None)
            response = json.loads(buf.split(b"\n")[0])
            return bool(response.get("ok"))
        except (BrokenPipeError, ConnectionResetError, OSError, json.JSONDecodeError, TimeoutError):
            return False

    def close(self) -> None:
        """Close the underlying socket."""
        try:
            self._sock.close()
        except OSError:
            pass
