from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any, Mapping

import requests

from acp.config import api_key
from acp.config import bearer_token as bearer_token_from_env
from acp.config import interceptor_base_url
from acp.exceptions import ACPSDKError, PolicyDenied, PolicyEscalated


@dataclass(frozen=True)
class ToolCallResponse:
    decision: str
    trace_id: str
    reason: str | None = None


@dataclass(frozen=True)
class ApprovalRecord:
    trace_id: str
    status: str
    agent_id: str
    tool: str
    action: str
    escalate_reason: str | None = None
    policy: str | None = None


class ACPClient:
    """HTTP client for the ACP interceptor `/tool-call` endpoint."""

    def __init__(
        self,
        base_url: str | None = None,
        timeout_s: float = 10.0,
        bearer_token: str | None = None,
    ) -> None:
        self._base = (base_url or interceptor_base_url()).rstrip("/")
        self._timeout_s = timeout_s
        tok = bearer_token if bearer_token is not None else bearer_token_from_env()
        self._bearer = tok.strip() if tok else None

    def tool_call(
        self,
        *,
        agent_id: str,
        tool: str,
        action: str,
        amount: int | None = None,
        risk_score: float | None = None,
        extra: Mapping[str, Any] | None = None,
        trace_id: str | None = None,
    ) -> ToolCallResponse:
        body: dict[str, Any] = {
            "agent_id": agent_id,
            "tool": tool,
            "action": action,
        }
        if amount is not None:
            body["amount"] = int(amount)
        if risk_score is not None:
            body["risk_score"] = float(risk_score)
        if extra:
            body["extra"] = dict(extra)

        headers = {"Content-Type": "application/json"}
        if trace_id:
            headers["X-Trace-ID"] = trace_id
        key = api_key()
        if key:
            headers["X-API-Key"] = key
        if self._bearer:
            headers["Authorization"] = f"Bearer {self._bearer}"

        url = f"{self._base}/tool-call"
        r = requests.post(url, json=body, headers=headers, timeout=self._timeout_s)
        r.raise_for_status()
        data = r.json()
        return ToolCallResponse(
            decision=str(data.get("decision", "")).lower(),
            trace_id=str(data.get("trace_id", "")),
            reason=data.get("reason"),
        )

    def enforce_or_raise(
        self,
        *,
        agent_id: str,
        tool: str,
        action: str,
        amount: int | None = None,
        risk_score: float | None = None,
        extra: Mapping[str, Any] | None = None,
        trace_id: str | None = None,
    ) -> ToolCallResponse:
        resp = self.tool_call(
            agent_id=agent_id,
            tool=tool,
            action=action,
            amount=amount,
            risk_score=risk_score,
            extra=extra,
            trace_id=trace_id,
        )
        if resp.decision == "allow":
            return resp
        payload = {
            "decision": resp.decision,
            "trace_id": resp.trace_id,
            "reason": resp.reason,
        }
        if resp.decision == "deny":
            raise PolicyDenied(
                trace_id=resp.trace_id,
                reason=resp.reason,
                decision_payload=payload,
            )
        if resp.decision == "escalate":
            raise PolicyEscalated(
                trace_id=resp.trace_id,
                reason=resp.reason,
                decision_payload=payload,
            )
        raise ACPSDKError(f"unknown_decision: {resp.decision!r} trace_id={resp.trace_id}")

    def _request_headers(self) -> dict[str, str]:
        headers: dict[str, str] = {}
        key = api_key()
        if key:
            headers["X-API-Key"] = key
        if self._bearer:
            headers["Authorization"] = f"Bearer {self._bearer}"
        return headers

    def get_approval(self, trace_id: str) -> ApprovalRecord:
        """Poll human approval state for a trace_id (escalation ticket)."""
        url = f"{self._base}/api/v1/approvals/{trace_id}"
        r = requests.get(url, headers=self._request_headers(), timeout=self._timeout_s)
        if r.status_code == 404:
            raise ACPSDKError(f"approval_not_found trace_id={trace_id}")
        r.raise_for_status()
        data = r.json()
        return ApprovalRecord(
            trace_id=str(data.get("trace_id", trace_id)),
            status=str(data.get("status", "")).lower(),
            agent_id=str(data.get("agent_id", "")),
            tool=str(data.get("tool", "")),
            action=str(data.get("action", "")),
            escalate_reason=data.get("escalate_reason"),
            policy=data.get("policy"),
        )

    def wait_for_approval(
        self,
        trace_id: str,
        *,
        timeout_s: float = 300.0,
        poll_s: float = 2.0,
    ) -> ApprovalRecord:
        """Block until trace_id is approved or denied (client-owned poll loop)."""
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            rec = self.get_approval(trace_id)
            if rec.status in ("approved", "denied"):
                return rec
            time.sleep(poll_s)
        raise ACPSDKError(f"approval_timeout trace_id={trace_id}")

    def enforce_after_approval_or_raise(
        self,
        *,
        agent_id: str,
        tool: str,
        action: str,
        trace_id: str,
        amount: int | None = None,
        risk_score: float | None = None,
        extra: Mapping[str, Any] | None = None,
        wait_timeout_s: float = 300.0,
        poll_s: float = 2.0,
    ) -> ToolCallResponse:
        """On escalate: poll until approved/denied, then retry tool-call with the same trace_id."""
        try:
            return self.enforce_or_raise(
                agent_id=agent_id,
                tool=tool,
                action=action,
                amount=amount,
                risk_score=risk_score,
                extra=extra,
                trace_id=trace_id,
            )
        except PolicyEscalated:
            rec = self.wait_for_approval(
                trace_id, timeout_s=wait_timeout_s, poll_s=poll_s
            )
            if rec.status == "denied":
                raise PolicyDenied(
                    trace_id=trace_id,
                    reason="approval_denied",
                    decision_payload={"decision": "deny", "trace_id": trace_id},
                )
            return self.enforce_or_raise(
                agent_id=agent_id,
                tool=tool,
                action=action,
                amount=amount,
                risk_score=risk_score,
                extra=extra,
                trace_id=trace_id,
            )
