"""Guard PDP client for Hermes plugins.

Talks to the local aarvion-guard daemon over a Unix socket (macOS/Linux) or
loopback TCP (Windows). Two auth factors: bearer token + kernel peer identity.
HMAC-SHA256 proof verification when a server key is configured.

This is a THIN client: it sends raw tool name + args to /v1/govern and lets
the guard's Go normalizer handle surface classification, HTTP semantic
extraction, and DLP scanning. No Python-side tool profiling.
"""
from __future__ import annotations

import hashlib
import hmac
import json
import logging
import os
import time
from typing import Any, TypedDict

import httpx

logger = logging.getLogger("aarvion.guard")

GUARD_CONFIG_PATH = os.path.expanduser("~/.aarvion/guard.json")

DEFAULT_TIMEOUT_S = 2.0
DEFAULT_APPROVAL_TIMEOUT_S = 90.0
DEFAULT_APPROVAL_POLL_S = 2.0
MAX_FIELD_BYTES = 16_384

SHELL_TOOLS = frozenset({"exec", "shell", "bash", "sh", "command", "run", "zsh"})

READONLY_TOOLS = frozenset({
    "read", "read_file", "file_read", "grep", "find", "ls", "list",
    "search", "glob", "web_search", "tool_search", "tool_describe",
    "get_goal", "sessions_list", "sessions_history", "session_status",
    "agents_list", "browser_screenshot",
})


class GuardVerdict(TypedDict, total=False):
    verdict: str  # "allow" | "deny" | "ask"
    reason: str
    policy_id: str
    decision_id: str
    fail_mode: bool


class GuardConfig(TypedDict, total=False):
    enabled: bool
    socket_path: str | None
    token: str | None
    server_key: str | None
    fail_mode: str  # "closed" | "open"
    timeout_s: float
    approval_timeout_s: float
    approval_poll_s: float


def _truncate(s: str) -> str:
    if len(s) <= MAX_FIELD_BYTES:
        return s
    return s[:MAX_FIELD_BYTES] + "…[truncated]"


def _bounded_args(args: dict[str, Any]) -> dict[str, Any]:
    out: dict[str, Any] = {}
    for k, v in args.items():
        out[k] = _truncate(v) if isinstance(v, str) else v
    return out


def _surface_for(tool_name: str) -> str:
    t = tool_name.lower()
    if t in SHELL_TOOLS:
        return "exec"
    if t.endswith("_send"):
        return "send"
    if t.startswith("browser_"):
        return "tool"
    if t in ("web_fetch", "http", "fetch", "http_request", "api_call"):
        return "egress"
    if t.startswith("file_"):
        return "tool"
    return "tool"


def _should_govern(tool_name: str) -> bool:
    return tool_name.lower() not in READONLY_TOOLS


def _shell_command(tool_name: str, args: dict[str, Any]) -> str | None:
    if tool_name.lower() not in SHELL_TOOLS:
        return None
    for k in ("command", "cmd", "input", "script"):
        if isinstance(args.get(k), str) and args[k]:
            return args[k]
    return None


def load_config() -> GuardConfig:
    """Load guard config from ~/.aarvion/guard.json."""
    if not os.path.exists(GUARD_CONFIG_PATH):
        return GuardConfig(enabled=False)
    with open(GUARD_CONFIG_PATH) as f:
        raw = json.load(f)
    gov = raw.get("govern", {})
    sock = gov.get("socket", {})
    fail_modes = gov.get("fail_modes", {})
    return GuardConfig(
        enabled=True,
        socket_path=sock.get("path"),
        token=sock.get("token"),
        server_key=sock.get("server_key"),
        fail_mode=fail_modes.get("send", "closed"),
        timeout_s=DEFAULT_TIMEOUT_S,
        approval_timeout_s=DEFAULT_APPROVAL_TIMEOUT_S,
        approval_poll_s=DEFAULT_APPROVAL_POLL_S,
    )


def _proof_valid(server_key: str, nonce: str, raw_body: bytes, proof: str | None) -> bool:
    if not proof:
        return False
    body_hash = hashlib.sha256(raw_body).hexdigest()
    want = hmac.new(
        server_key.encode(), f"{nonce}.{body_hash}".encode(), hashlib.sha256
    ).digest()
    got = bytes.fromhex(proof)
    return hmac.compare_digest(got, want)


def _make_client(config: GuardConfig) -> httpx.Client:
    socket_path = config.get("socket_path")
    if socket_path:
        transport = httpx.HTTPTransport(uds=socket_path)
        return httpx.Client(transport=transport, base_url="http://guard")
    return httpx.Client(base_url="http://127.0.0.1:9119")


def _build_govern_body(
    tool_name: str,
    args: dict[str, Any],
    agent_id: str = "",
    session_id: str = "",
) -> tuple[str, str]:
    nonce = os.urandom(16).hex()
    cmd = _shell_command(tool_name, args)
    payload = json.dumps({
        "contract_version": "1",
        "nonce": nonce,
        "ctx": {
            "surface": _surface_for(tool_name),
            "phase": "pre",
            "caller": {
                "principal_id": agent_id,
                "session_id": session_id,
                "source": "hermes",
                "tool": tool_name,
                "trust": "untrusted",
            },
        },
        "action": {
            "tool": tool_name,
            "operation": "exec" if cmd is not None else "call",
            "args": {**_bounded_args(args), **({"cmd": cmd} if cmd else {})},
        },
        "attributes": {
            "request": {
                "http": {
                    "method": "",
                    "host": "",
                    "path": "",
                    "body": "",
                    "headers": {},
                },
            },
        },
    })
    return payload, nonce


def request_verdict(
    tool_name: str,
    args: dict[str, Any],
    config: GuardConfig,
    agent_id: str = "",
    session_id: str = "",
) -> GuardVerdict | None:
    """POST /v1/govern and return the parsed verdict, or None on transport failure."""
    payload, nonce = _build_govern_body(tool_name, args, agent_id, session_id)
    try:
        client = _make_client(config)
        resp = client.post(
            "/v1/govern",
            content=payload,
            headers={
                "content-type": "application/json",
                "authorization": f"Bearer {config.get('token', '')}",
            },
            timeout=config.get("timeout_s", DEFAULT_TIMEOUT_S),
        )
        raw_body = resp.content
        server_key = config.get("server_key")
        if server_key:
            proof = resp.headers.get("x-aarvion-guard-proof")
            if not _proof_valid(server_key, nonce, raw_body, proof):
                logger.warning("Guard proof verification failed")
                return None
        data = resp.json()
        v = data.get("verdict")
        if v not in ("allow", "deny", "ask"):
            return None
        return GuardVerdict(
            verdict=v,
            reason=data.get("reason", ""),
            policy_id=data.get("policy_id", ""),
            decision_id=data.get("decision_id", ""),
        )
    except Exception:
        logger.debug("Guard request failed", exc_info=True)
        return None


def fetch_approval_status(
    decision_id: str,
    config: GuardConfig,
) -> str:
    """GET /v1/approvals/{id}. Returns 'pending', 'allow', 'deny', or 'unreachable'."""
    try:
        client = _make_client(config)
        nonce = os.urandom(16).hex()
        resp = client.get(
            f"/v1/approvals/{decision_id}",
            params={"nonce": nonce},
            headers={"authorization": f"Bearer {config.get('token', '')}"},
            timeout=config.get("timeout_s", DEFAULT_TIMEOUT_S),
        )
        if resp.status_code < 200 or resp.status_code >= 300:
            return "unreachable"
        raw_body = resp.content
        server_key = config.get("server_key")
        if server_key:
            proof = resp.headers.get("x-aarvion-guard-proof")
            if not _proof_valid(server_key, nonce, raw_body, proof):
                return "unreachable"
        data = resp.json()
        status = data.get("status", "")
        if status in ("pending", "allow", "deny"):
            return status
        return "unreachable"
    except Exception:
        return "unreachable"


def await_approval(
    decision_id: str,
    config: GuardConfig,
) -> str:
    """Poll until the human resolves the approval. Returns 'allow' or 'deny'."""
    if not decision_id:
        return "deny"
    timeout = config.get("approval_timeout_s", DEFAULT_APPROVAL_TIMEOUT_S)
    poll_interval = config.get("approval_poll_s", DEFAULT_APPROVAL_POLL_S)
    deadline = time.monotonic() + timeout
    while True:
        status = fetch_approval_status(decision_id, config)
        if status == "allow":
            return "allow"
        if status == "deny":
            return "deny"
        if time.monotonic() >= deadline:
            return "deny"
        time.sleep(poll_interval)
        if time.monotonic() >= deadline:
            return "deny"


def _build_inbound_body(
    channel: str,
    sender: str,
    message: str,
    agent_id: str = "",
    session_id: str = "",
) -> tuple[str, str]:
    nonce = os.urandom(16).hex()
    payload = json.dumps({
        "contract_version": "1",
        "nonce": nonce,
        "ctx": {
            "surface": "inbound",
            "phase": "pre",
            "caller": {
                "principal_id": agent_id,
                "session_id": session_id,
                "source": "hermes",
                "tool": "gateway_dispatch",
                "trust": "untrusted",
            },
        },
        "action": {
            "tool": "gateway_dispatch",
            "operation": "receive",
            "args": {
                "channel": channel,
                "sender": sender,
                "body": _truncate(message),
            },
        },
        "attributes": {
            "request": {
                "http": {
                    "method": "POST",
                    "host": channel,
                    "path": "/inbound",
                    "body": _truncate(message),
                    "headers": {},
                },
            },
        },
    })
    return payload, nonce


def request_inbound_verdict(
    channel: str,
    sender: str,
    message: str,
    config: GuardConfig,
    agent_id: str = "",
    session_id: str = "",
) -> GuardVerdict | None:
    payload, nonce = _build_inbound_body(channel, sender, message, agent_id, session_id)
    try:
        client = _make_client(config)
        resp = client.post(
            "/v1/govern",
            content=payload,
            headers={
                "content-type": "application/json",
                "authorization": f"Bearer {config.get('token', '')}",
            },
            timeout=config.get("timeout_s", DEFAULT_TIMEOUT_S),
        )
        raw_body = resp.content
        server_key = config.get("server_key")
        if server_key:
            proof = resp.headers.get("x-aarvion-guard-proof")
            if not _proof_valid(server_key, nonce, raw_body, proof):
                logger.warning("Guard proof verification failed")
                return None
        data = resp.json()
        v = data.get("verdict")
        if v not in ("allow", "deny", "ask"):
            return None
        return GuardVerdict(
            verdict=v,
            reason=data.get("reason", ""),
            policy_id=data.get("policy_id", ""),
            decision_id=data.get("decision_id", ""),
        )
    except Exception:
        logger.debug("Guard inbound request failed", exc_info=True)
        return None


def evaluate(
    tool_name: str,
    args: dict[str, Any],
    config: GuardConfig | None = None,
    agent_id: str = "",
    session_id: str = "",
) -> GuardVerdict:
    """High-level entry point. Returns a concrete verdict, never throws.

    - Not enabled or read-only tool → allow
    - Guard reachable → policy verdict (with approval polling for 'ask')
    - Guard unreachable → fail-mode (default: deny)
    """
    if config is None:
        config = load_config()
    if not config.get("enabled"):
        return GuardVerdict(verdict="allow")
    if not _should_govern(tool_name):
        return GuardVerdict(verdict="allow")

    verdict = request_verdict(tool_name, args, config, agent_id, session_id)
    if verdict is None:
        fail = config.get("fail_mode", "closed")
        if fail == "open":
            return GuardVerdict(verdict="allow", fail_mode=True, reason="guard_unreachable_fail_open")
        return GuardVerdict(verdict="deny", fail_mode=True, reason="guard_unreachable_fail_closed")

    if verdict["verdict"] == "ask":
        decision_id = verdict.get("decision_id", "")
        result = await_approval(decision_id, config)
        if result != "allow":
            return GuardVerdict(verdict="deny", reason="approval_denied_or_timeout")
        return GuardVerdict(verdict="allow", reason="approved")

    return verdict


def evaluate_inbound(
    channel: str,
    sender: str,
    message: str,
    config: GuardConfig | None = None,
    agent_id: str = "",
    session_id: str = "",
) -> GuardVerdict:
    if config is None:
        config = load_config()
    if not config.get("enabled"):
        return GuardVerdict(verdict="allow")

    verdict = request_inbound_verdict(channel, sender, message, config, agent_id, session_id)
    if verdict is None:
        fail = config.get("fail_mode", "closed")
        if fail == "open":
            return GuardVerdict(verdict="allow", fail_mode=True, reason="guard_unreachable_fail_open")
        return GuardVerdict(verdict="deny", fail_mode=True, reason="guard_unreachable_fail_closed")

    return verdict
