from __future__ import annotations

import json
import socket
import time
from pathlib import Path
from typing import Any
from uuid import UUID, uuid4

from .redaction import redact_text


class AgentError(RuntimeError):
    pass


class AgentClient:
    PROTOCOL_VERSION = 1
    MAX_MESSAGE_BYTES = 64 * 1024

    def __init__(self, socket_path: Path, *, timeout: int = 30):
        self.socket_path = Path(socket_path)
        self.timeout = timeout

    def request(self, method: str, path: str, **values: Any) -> dict[str, Any]:
        request = {
            "protocol_version": self.PROTOCOL_VERSION,
            "method": method,
            "path": path,
            **values,
        }
        encoded = json.dumps(request, separators=(",", ":"), ensure_ascii=True).encode("utf-8") + b"\n"
        if len(encoded) > self.MAX_MESSAGE_BYTES:
            raise AgentError("Agent request exceeds 64 KiB.")
        if not hasattr(socket, "AF_UNIX"):
            raise AgentError("Unix sockets are unavailable on this platform.")
        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
                client.settimeout(self.timeout)
                client.connect(str(self.socket_path))
                client.sendall(encoded)
                client.shutdown(socket.SHUT_WR)
                response = bytearray()
                while b"\n" not in response:
                    chunk = client.recv(4096)
                    if not chunk:
                        break
                    response.extend(chunk)
                    if len(response) > self.MAX_MESSAGE_BYTES:
                        raise AgentError("Agent response exceeds 64 KiB.")
        except (OSError, TimeoutError) as exc:
            raise AgentError(f"Agent connection failed: {type(exc).__name__}") from exc
        line, separator, remainder = bytes(response).partition(b"\n")
        if not separator or remainder:
            raise AgentError("Agent must return exactly one JSON line.")
        try:
            payload = json.loads(line.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise AgentError("Agent returned invalid JSON.") from exc
        if not isinstance(payload, dict) or payload.get("protocol_version") != self.PROTOCOL_VERSION:
            raise AgentError("Agent returned an incompatible protocol response.")
        status = payload.get("status")
        body = payload.get("body")
        if not isinstance(status, int) or not isinstance(body, dict):
            raise AgentError("Agent returned a malformed response.")
        if status < 200 or status >= 300:
            message = body.get("error") or body.get("detail") or f"Agent returned status {status}."
            raise AgentError(redact_text(message, limit=2_000))
        return body

    def capabilities(self) -> dict[str, Any]:
        return self.request("GET", "/v1/capabilities")

    def submit_operation(
        self,
        *,
        request_id: str,
        action: str,
        plugin_slug: str,
        version: str | None,
        approved_payload_sha256: str | None,
        requested_by: str,
    ) -> UUID:
        idempotency_key = uuid4()
        body = self.request(
            "POST",
            "/v1/operations",
            idempotency_key=str(idempotency_key),
            body={
                "request_id": request_id,
                "action": action,
                "plugin_slug": plugin_slug,
                "version": version,
                "approved_payload_sha256": approved_payload_sha256,
                "requested_by": requested_by,
            },
        )
        raw_id = body.get("operation_id") or body.get("id")
        try:
            return UUID(str(raw_id))
        except (ValueError, TypeError) as exc:
            raise AgentError("Agent did not return a valid operation UUID.") from exc

    def wait_for_operation(
        self,
        operation_id: UUID,
        *,
        overall_timeout: int,
        poll_interval: float,
    ) -> dict[str, Any]:
        deadline = time.monotonic() + overall_timeout
        while True:
            body = self.get_operation(operation_id)
            state = body.get("state") or body.get("status")
            if state in {"succeeded", "completed"}:
                result = body.get("result", body)
                if not isinstance(result, dict):
                    raise AgentError("Agent operation result is malformed.")
                return result
            if state in {"failed", "errored", "cancelled"}:
                raise AgentError(redact_text(body.get("error") or f"Agent operation {state}."))
            if state not in {"queued", "pending", "running", "accepted"}:
                raise AgentError("Agent returned an unknown operation state.")
            if time.monotonic() >= deadline:
                raise AgentError("Timed out waiting for the host agent operation.")
            time.sleep(poll_interval)

    def get_operation(self, operation_id: UUID) -> dict[str, Any]:
        return self.request("GET", f"/v1/operations/{operation_id}")