"""Cross-platform local socket client.

Unix domain sockets on Linux/macOS, named pipes on Windows.
Thread-safe: each operation holds a lock for the full request/response cycle.
Auto-reconnects on daemon restart.
"""

import json
import os
import platform
import socket
import threading
from typing import Optional

MAX_PAYLOAD = 256 * 1024 * 1024  # 256 MB guard


def default_socket_addr() -> str:
    """Platform-appropriate socket address matching the Rust daemon."""
    override = os.environ.get("IROHDS_SOCKET")
    if override:
        return override
    if platform.system() == "Windows":
        return r"\\.\pipe\irohds"
    runtime = os.environ.get("XDG_RUNTIME_DIR", "/tmp")
    return os.path.join(runtime, "irohds.sock")


class DaemonClient:
    def __init__(self, address: Optional[str] = None):
        self._addr = address or default_socket_addr()
        self._sock: Optional[socket.socket] = None
        self._file = None
        self._lock = threading.Lock()

    def _connect(self):
        if self._sock is not None:
            return
        if platform.system() == "Windows":
            try:
                self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                self._sock.connect(self._addr)
            except (AttributeError, OSError):
                self._file = open(self._addr, "r+b", buffering=0)
                return
        else:
            self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self._sock.connect(self._addr)
        self._file = self._sock.makefile("rb")

    def _disconnect(self):
        for obj in (self._file, self._sock):
            if obj:
                try: obj.close()
                except OSError: pass
        self._sock = None
        self._file = None

    def _reconnect(self):
        self._disconnect()
        self._connect()

    def _send(self, req: dict):
        data = (json.dumps(req) + "\n").encode()
        if self._sock:
            self._sock.sendall(data)
        elif self._file:
            self._file.write(data)

    def _recv_line(self) -> dict:
        line = self._file.readline()
        if not line:
            raise ConnectionError("daemon closed")
        return json.loads(line)

    def _recv_bytes(self, n: int) -> bytes:
        if n > MAX_PAYLOAD:
            raise ValueError(f"payload too large: {n}")
        buf = b""
        while len(buf) < n:
            chunk = self._file.read(n - len(buf))
            if not chunk:
                raise ConnectionError("short read")
            buf += chunk
        return buf

    def _request(self, req: dict):
        """Send request with one retry on connection failure."""
        try:
            self._connect()
            self._send(req)
        except (OSError, ConnectionError):
            self._reconnect()
            self._send(req)

    def ping(self) -> bool:
        with self._lock:
            try:
                self._request({"op": "ping"})
                r = self._recv_line()
                return r.get("status") == "pong"
            except Exception:
                self._disconnect()
                return False

    def probe(self, ns: str, key: str) -> Optional[dict]:
        """Metadata-only lookup. Returns FnCost fields if hit, None if miss."""
        with self._lock:
            self._request({"op": "probe", "ns": ns, "key": key})
            r = self._recv_line()
            if r["status"] == "probe_hit":
                return {
                    "ops": r["ops"],
                    "mem_bytes": r["mem_bytes"],
                    "peak_mem": r["peak_mem"],
                    "result_bytes": r["size"],
                }
            if r["status"] == "miss":
                return None
            if r["status"] == "error":
                raise RuntimeError(r.get("message"))
            raise RuntimeError(f"unexpected: {r}")

    def get(self, ns: str, key: str) -> Optional[bytes]:
        with self._lock:
            self._request({"op": "get", "ns": ns, "key": key})
            r = self._recv_line()
            if r["status"] == "hit":
                return self._recv_bytes(r["size"])
            if r["status"] == "miss":
                return None
            if r["status"] == "error":
                raise RuntimeError(r.get("message"))
            raise RuntimeError(f"unexpected: {r}")

    def check(self, ns: str, key: str) -> Optional[bytes]:
        """Combined probe+get in one IPC round-trip. Local data only.

        Returns cached bytes if present locally, None if miss.
        Does not trigger network fetch (use get() for that).
        """
        with self._lock:
            self._request({"op": "check", "ns": ns, "key": key})
            r = self._recv_line()
            if r["status"] == "hit":
                return self._recv_bytes(r["size"])
            if r["status"] == "miss":
                return None
            if r["status"] == "error":
                raise RuntimeError(r.get("message"))
            raise RuntimeError(f"unexpected: {r}")

    def put(self, ns: str, key: str, data: bytes, compute_ns: int = 0,
            ops: int = 0, mem_bytes: int = 0, peak_mem: int = 0, tag: str = ""):
        meta = {}
        if compute_ns: meta["compute_ns"] = compute_ns
        if ops: meta["ops"] = ops
        if mem_bytes: meta["mem_bytes"] = mem_bytes
        if peak_mem: meta["peak_mem"] = peak_mem
        if tag: meta["tag"] = tag
        with self._lock:
            self._request({"op": "put", "ns": ns, "key": key, "size": len(data), "meta": meta})
            if self._sock:
                self._sock.sendall(data)
            elif self._file:
                self._file.write(data)
            r = self._recv_line()
            if r.get("status") == "error":
                raise RuntimeError(r.get("message"))

    def join(self, ns: str):
        with self._lock:
            self._request({"op": "join", "ns": ns})
            r = self._recv_line()
            if r.get("status") == "error":
                raise RuntimeError(r.get("message"))

    def status(self, ns: str) -> dict:
        with self._lock:
            self._request({"op": "status", "ns": ns})
            return self._recv_line()

    def evict(self, ns: str, tag: str):
        with self._lock:
            self._request({"op": "evict", "ns": ns, "tag": tag})
            return self._recv_line()

    def put_file(self, ns: str, key: str, path: str,
                 compute_ns: int = 0, ops: int = 0, mem_bytes: int = 0,
                 peak_mem: int = 0, tag: str = "") -> dict:
        """Import a file into the blob store as a chunked HashSeq."""
        meta = {}
        if compute_ns: meta["compute_ns"] = compute_ns
        if ops: meta["ops"] = ops
        if mem_bytes: meta["mem_bytes"] = mem_bytes
        if peak_mem: meta["peak_mem"] = peak_mem
        if tag: meta["tag"] = tag
        with self._lock:
            self._request({"op": "put_file", "ns": ns, "key": key,
                           "path": path, "meta": meta})
            r = self._recv_line()
            if r.get("status") == "error":
                raise RuntimeError(r.get("message"))
            return r

    def get_file(self, ns: str, blob_hash: str, path: str) -> Optional[int]:
        """Materialize a HashSeq to a file on disk."""
        with self._lock:
            self._request({"op": "get_file", "ns": ns, "hash": blob_hash, "path": path})
            r = self._recv_line()
            if r["status"] == "hit":
                return r["size"]
            if r["status"] == "miss":
                return None
            if r["status"] == "error":
                raise RuntimeError(r.get("message"))
            raise RuntimeError(f"unexpected: {r}")

    def close(self):
        with self._lock:
            self._disconnect()
