"""
Agent Protocol — The standard API contract for all local Brynq agents.

Every agent that runs locally in the sandbox MUST expose these HTTP
endpoints on the port assigned by the sandbox manager:

    POST /execute  — Run a task (prompt + context in, output out)
    GET  /health   — Liveness probe (200 = alive)
    GET  /info     — Agent metadata (name, version, capabilities)

This module defines the request/response schemas as frozen dataclasses
(no external dependencies) and a small client that the runtime uses to
talk to a running local agent.

Design:
    - JSON over HTTP, no streaming (agents are short-lived workers).
    - Timeouts enforced by the caller (sandbox manager), not the agent.
    - The agent MUST NOT call any endpoint outside its declared whitelist.
"""

from __future__ import annotations

import json
import logging
import time
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass, field
from typing import Any

log = logging.getLogger("brynq.runtime.agent_protocol")

AGENT_PROTOCOL_VERSION: int = 1

# ── Request / Response Schemas ──────────────────────────────────────────


@dataclass(frozen=True)
class AgentExecuteRequest:
    """Payload sent to ``POST /execute``.

    Attributes:
        prompt: The task instruction for the agent.
        context: Upstream outputs from earlier chain steps. May be empty.
        config: Arbitrary key/value config the agent can use (e.g.
                temperature, max_tokens, output_format).
    """

    prompt: str
    context: dict[str, Any] = field(default_factory=dict)
    config: dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps(asdict(self))


@dataclass(frozen=True)
class AgentExecuteResponse:
    """Payload returned by ``POST /execute``.

    Attributes:
        output: The agent's result text.
        tokens_used: Self-reported token consumption (0 if unknown).
        elapsed_ms: Wall-clock time the agent spent processing.
        metadata: Optional extra info (citations, confidence, etc.).
    """

    output: str
    tokens_used: int = 0
    elapsed_ms: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, d: dict) -> AgentExecuteResponse:
        return cls(
            output=d.get("output", ""),
            tokens_used=int(d.get("tokens_used", 0)),
            elapsed_ms=float(d.get("elapsed_ms", 0.0)),
            metadata=d.get("metadata", {}),
        )


@dataclass(frozen=True)
class AgentHealthResponse:
    """Payload returned by ``GET /health``.

    Attributes:
        status: "ok" when the agent is healthy.
        uptime_seconds: How long the agent process has been alive.
    """

    status: str = "ok"
    uptime_seconds: float = 0.0

    @classmethod
    def from_dict(cls, d: dict) -> AgentHealthResponse:
        return cls(
            status=d.get("status", "unknown"),
            uptime_seconds=float(d.get("uptime_seconds", 0.0)),
        )


@dataclass(frozen=True)
class AgentInfoResponse:
    """Payload returned by ``GET /info``.

    Attributes:
        name: Human-readable agent name.
        version: Semantic version string.
        capabilities: List of capability tags (e.g. ["summarization", "code"]).
        protocol_version: Protocol version the agent implements.
    """

    name: str = ""
    version: str = "0.0.0"
    capabilities: list[str] = field(default_factory=list)
    protocol_version: int = AGENT_PROTOCOL_VERSION

    @classmethod
    def from_dict(cls, d: dict) -> AgentInfoResponse:
        return cls(
            name=d.get("name", ""),
            version=d.get("version", "0.0.0"),
            capabilities=d.get("capabilities", []),
            protocol_version=int(d.get("protocol_version", AGENT_PROTOCOL_VERSION)),
        )


# ── Protocol Client ─────────────────────────────────────────────────────


class AgentProtocolClient:
    """Talks to a running local agent via its sandboxed HTTP endpoint.

    The client enforces a request timeout but does NOT enforce sandbox
    rules — that is the sandbox manager's job (network restrictions,
    filesystem isolation, etc.).

    Attributes:
        base_url: Root URL of the agent (e.g. ``http://127.0.0.1:9100``).
        timeout_seconds: Per-request timeout.
    """

    def __init__(self, base_url: str, timeout_seconds: float = 30.0) -> None:
        # Strip trailing slash for clean URL construction.
        self.base_url = base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds

    # ── Public API ───────────────────────────────────────────────────

    def execute(self, request: AgentExecuteRequest) -> AgentExecuteResponse:
        """POST /execute — run a task on the agent.

        Args:
            request: The execute payload.

        Returns:
            AgentExecuteResponse with the agent's output.

        Raises:
            AgentProtocolError: On HTTP error, timeout, or malformed response.
        """
        body = request.to_json().encode("utf-8")
        resp = self._post("/execute", body)
        return AgentExecuteResponse.from_dict(resp)

    def health(self) -> AgentHealthResponse:
        """GET /health — check agent liveness.

        Returns:
            AgentHealthResponse (status should be "ok").

        Raises:
            AgentProtocolError: If the agent is unreachable or unhealthy.
        """
        resp = self._get("/health")
        return AgentHealthResponse.from_dict(resp)

    def info(self) -> AgentInfoResponse:
        """GET /info — retrieve agent metadata.

        Returns:
            AgentInfoResponse with name, version, capabilities.

        Raises:
            AgentProtocolError: On connection or parse errors.
        """
        resp = self._get("/info")
        return AgentInfoResponse.from_dict(resp)

    # ── Internals ────────────────────────────────────────────────────

    def _get(self, path: str) -> dict:
        url = f"{self.base_url}{path}"
        req = urllib.request.Request(url, method="GET")
        req.add_header("Accept", "application/json")
        return self._send(req)

    def _post(self, path: str, body: bytes) -> dict:
        url = f"{self.base_url}{path}"
        req = urllib.request.Request(url, data=body, method="POST")
        req.add_header("Content-Type", "application/json")
        req.add_header("Accept", "application/json")
        return self._send(req)

    def _send(self, req: urllib.request.Request) -> dict:
        try:
            with urllib.request.urlopen(req, timeout=self.timeout_seconds) as resp:
                raw = resp.read().decode("utf-8")
                return json.loads(raw) if raw else {}
        except urllib.error.HTTPError as exc:
            raise AgentProtocolError(
                f"Agent returned HTTP {exc.code}: {exc.reason}"
            ) from exc
        except urllib.error.URLError as exc:
            raise AgentProtocolError(
                f"Cannot reach agent at {req.full_url}: {exc.reason}"
            ) from exc
        except json.JSONDecodeError as exc:
            raise AgentProtocolError(
                f"Agent returned invalid JSON: {exc}"
            ) from exc
        except TimeoutError as exc:
            raise AgentProtocolError(
                f"Agent request timed out after {self.timeout_seconds}s"
            ) from exc


class AgentProtocolError(Exception):
    """Raised when communication with a local agent fails."""
