"""
Python client for Squidbrake. Drop this file into any agent / service.

    from client import Gateway, Denied

    gw = Gateway("https://gateway.example.com", api_key="...", source="my-agent", session_id=run_id)

    @gw.guard()                       # every call is checked + recorded; raises Denied if blocked
    def shell_exec(command: str) -> str:
        ...

    @gw.guard(name="web.fetch")       # works on async functions too
    async def fetch(url: str) -> str:
        ...

    gw.log("deploy", input={"env": "prod"}, output="ok", kind="action")   # record-only

    @gw.guard(snapshot=lambda args: db.get_order(args["order_id"]))   # undo: keep the record before it changes
    def update_order(order_id: str, status: str): ...

    gw.usage(model="claude-opus-5-5", input_tokens=1200, output_tokens=300)   # for rules.yaml `budgets:`

Calls matched by a `review` rule block inside check() / the decorator until a human approves
(the tool then runs) or rejects / the approval times out (Denied is raised).
"""
from __future__ import annotations

import asyncio
import functools
import inspect
import logging
import os
import time
from typing import Any, Callable

import httpx

try:  # what to do while the gateway is down (rules.yaml when_unreachable); optional when this file is copied alone
    import offline
except Exception:  # pragma: no cover
    offline = None

log = logging.getLogger("gateway.client")
LONG_POLL = 25.0  # seconds per /decision request while waiting for a human


class Denied(PermissionError):
    def __init__(self, name: str, reason: str, event_id: str | None, decided_by: str | None = None):
        by = f" (rejected by {decided_by})" if decided_by and decided_by != "timeout" else \
             " (approval timed out)" if decided_by == "timeout" else ""
        super().__init__(f"Gateway denied '{name}': {reason}{by}")
        self.name, self.reason, self.event_id, self.decided_by = name, reason, event_id, decided_by


class GatewayUnavailable(RuntimeError):
    pass


def refusal(e: Exception) -> str:
    """What an agent should read when Squidbrake stopped its tool call (returned as the tool's result by guard_tools,
    because some frameworks end the whole run on a tool exception instead of telling the model)."""
    if isinstance(e, Denied) and e.decided_by and e.decided_by != "timeout":
        return f"NOT RUN: a person rejected {e.name} in Squidbrake ({e.reason}). Don't retry it; ask the user how to proceed."
    if isinstance(e, Denied):
        return f"NOT RUN: Squidbrake stopped {e.name}: {e.reason}. Don't try to work around it; tell the user."
    return f"NOT RUN: {e}. Every action must go through Squidbrake, so nothing was done. Tell the user."


def _log_pending(name: str, decision: dict) -> None:
    log.warning("'%s' is waiting for human approval (event %s, expires %s): %s",
                name, decision["event_id"], decision.get("approval_deadline"), decision["reason"])


class Gateway:
    def __init__(self, url: str | None = None, api_key: str | None = None, *, source: str | None = None,
                 session_id: str | None = None, fail_open: bool = False, timeout: float = 5.0,
                 on_approval_pending: Callable[[str, dict], None] | None = _log_pending,
                 on_note: Callable[[dict], None] | None = None):
        """fail_open=False (default): if the gateway is unreachable, tools do NOT run.
        fail_open=True: tools run anyway and the call goes unrecorded. (A call already held for
        approval never runs without a decision, even with fail_open.)
        on_approval_pending(name, decision) is called once when a call is held for a human.
        on_note({by, text, at}) gets each note a person sends this agent (put it in front of your model); without it,
        notes wait for an agent that can pass them on."""
        self.url = (url or os.environ["GATEWAY_URL"]).rstrip("/")
        self.source, self.session_id, self.fail_open = source, session_id, fail_open
        self.on_approval_pending, self.on_note = on_approval_pending, on_note
        headers = {"X-Gateway-Key": api_key or os.getenv("GATEWAY_API_KEY", "")}
        self._http = httpx.Client(base_url=self.url, headers=headers, timeout=timeout)
        self._headers, self._timeout = headers, timeout
        self._aclients: dict = {}   # one async client per event loop: a client can't outlive the loop it was made in

    def _ahttp(self) -> httpx.AsyncClient:
        loop = asyncio.get_running_loop()
        for gone in [l for l in self._aclients if l.is_closed()]:
            self._aclients.pop(gone)
        if loop not in self._aclients:
            self._aclients[loop] = httpx.AsyncClient(base_url=self.url, headers=self._headers, timeout=self._timeout)
        return self._aclients[loop]

    # ---- payload helpers
    def _event(self, name, input, kind, metadata, **extra) -> dict:
        metadata = {**({"accepts_notes": True} if self.on_note else {}), **(metadata or {})}
        return {"name": name, "kind": kind, "input": input, "source": self.source,
                "session_id": self.session_id, "metadata": metadata, **extra}

    def _first(self, name: str, resp: httpx.Response | None, input: Any = None, kind: str = "tool_call") -> dict | None:
        if resp is None:
            if self.fail_open:
                return None
            if offline is not None:
                run, why = offline.decide(self.url, name=name, kind=kind, source=self.source,
                                          session_id=self.session_id, input=input)
                if run:
                    log.warning("Squidbrake unreachable; '%s' ran unrecorded: %s", name, why)
                    return None
                raise GatewayUnavailable(f"gateway unreachable, refusing to run '{name}': {why}")
            raise GatewayUnavailable(f"gateway unreachable, refusing to run '{name}'")
        resp.raise_for_status()
        d = resp.json()
        for n in d.get("notes") or []:
            self.on_note(n)
        if d["decision"] == "review" and self.on_approval_pending:
            self.on_approval_pending(name, d)
        return d

    @staticmethod
    def _final(name: str, d: dict | None) -> str | None:
        if d is None:
            return None
        if d["decision"] != "allow":
            raise Denied(name, d.get("decision_note") or d["reason"], d["event_id"], d.get("decided_by"))
        return d["event_id"]

    @staticmethod
    def _gave_up(d: dict) -> bool:
        # The server decides by the deadline; allow generous slack for clock skew / outages.
        from datetime import datetime, timezone
        deadline = datetime.fromisoformat(d["approval_deadline"].replace("Z", "+00:00"))
        return (datetime.now(timezone.utc) - deadline).total_seconds() > 60

    # ---- sync API
    def _snapshot_body(self, name: str, d: dict, input: Any, snapshot: Callable | None) -> dict | None:
        """What to tell the gateway about the undo copy for this call (None: no copy wanted / possible)."""
        if not d.get("snapshot"):
            return None
        if snapshot is None:
            return {"skipped": f"no snapshot function was given for {name}"} if d.get("snapshot_strict") else None
        try:
            data = snapshot(input)
        except Exception as e:
            return {"skipped": f"the snapshot function failed ({type(e).__name__}: {e})"}
        return {"kind": "data", "summary": f"{name}: the record before the change", "data": data}

    def _after_snapshot(self, name: str, d: dict, body: dict | None, posted: bool) -> None:
        if body and body.get("skipped") and d.get("snapshot") == "required":
            raise Denied(name, f"no undo copy could be made ({body['skipped']}), so it didn't run", d["event_id"])
        if body and not body.get("skipped") and not posted and d.get("snapshot") == "required":
            raise Denied(name, "the undo copy couldn't be saved on the gateway, so it didn't run", d["event_id"])

    def _keep(self, name: str, d: dict | None, input: Any, snapshot: Callable | None) -> None:
        if d is None or d.get("decision") != "allow":
            return
        body = self._snapshot_body(name, d, input, snapshot)
        posted = False
        if body:
            try:
                posted = self._http.post(f"/v1/events/{d['event_id']}/snapshot", json=body).status_code == 200
            except httpx.TransportError:
                pass
            if (body.get("skipped") or not posted) and d.get("snapshot") == "required":
                self.result(d["event_id"], error=f"not run: no undo copy ({body.get('skipped') or 'not saved'})")
        self._after_snapshot(name, d, body, posted)

    async def _akeep(self, name: str, d: dict | None, input: Any, snapshot: Callable | None) -> None:
        if d is None or d.get("decision") != "allow":
            return
        body = self._snapshot_body(name, d, input, snapshot)
        posted = False
        if body:
            try:
                posted = (await self._ahttp().post(f"/v1/events/{d['event_id']}/snapshot", json=body)).status_code == 200
            except httpx.TransportError:
                pass
            if (body.get("skipped") or not posted) and d.get("snapshot") == "required":
                await self.aresult(d["event_id"], error=f"not run: no undo copy ({body.get('skipped') or 'not saved'})")
        self._after_snapshot(name, d, body, posted)

    def usage(self, *, model: str | None = None, input_tokens: int = 0, output_tokens: int = 0,
              cost_usd: float | None = None) -> dict | None:
        """Report model tokens this agent spent, for rules.yaml `budgets:`. Never raises."""
        try:
            return self._http.post("/v1/usage", json={"source": self.source, "session_id": self.session_id,
                                                       "model": model, "input_tokens": input_tokens,
                                                       "output_tokens": output_tokens, "cost_usd": cost_usd}).json()
        except (httpx.HTTPError, ValueError):
            return None

    def check(self, name: str, input: Any = None, *, kind: str = "tool_call", metadata: dict | None = None,
              snapshot: Callable[[Any], Any] | None = None) -> str | None:
        """Ask before running. Returns event_id; raises Denied if blocked. Waits while a human reviews it.
        snapshot(input) returns the record as it is now; when the rules want an undo copy it's kept on the gateway
        before the call runs, and if it fails where a copy is required, the call doesn't run (Denied)."""
        d = self._decided(name, input, kind, metadata)
        eid = self._final(name, d)
        self._keep(name, d, input, snapshot)
        return eid

    def ask(self, question: str) -> str:
        """Ask a person on the team and wait for the answer (they answer in the dashboard or on their phone).
        Raises Denied if they decline, nobody answers in time, or the gateway is unreachable and won't take it."""
        d = self._decided("ask_user", {"question": question}, "question", None)
        if d is None:
            raise GatewayUnavailable("Squidbrake is unreachable, so nobody could be asked")
        self._final("ask_user", d)
        self.result(d["event_id"], output={"answer": d.get("decision_note")})
        return d.get("decision_note") or ""

    def _decided(self, name: str, input: Any, kind: str, metadata: dict | None) -> dict | None:
        """Send the call and wait while a person decides. None: unreachable and the rules let it run unrecorded."""
        try:
            resp = self._http.post("/v1/events", json=self._event(name, input, kind, metadata))
        except httpx.TransportError:
            resp = None
        d = self._first(name, resp, input, kind)
        while d is not None and d["decision"] == "review":
            try:
                r = self._http.get(f"/v1/events/{d['event_id']}/decision",
                                   params={"wait": LONG_POLL}, timeout=LONG_POLL + 10)
                r.raise_for_status()
                d = r.json()
            except httpx.TransportError:
                if self._gave_up(d):
                    raise GatewayUnavailable(f"lost the gateway while '{name}' awaited approval")
                time.sleep(2)
        return d

    def result(self, event_id: str | None, *, output: Any = None, error: str | None = None,
               duration_ms: float | None = None) -> None:
        if event_id is None:
            return
        try:
            self._http.post(f"/v1/events/{event_id}/result",
                            json={"output": output, "error": error, "duration_ms": duration_ms})
        except httpx.TransportError:
            pass  # the call already ran; don't fail the tool because reporting failed

    def log(self, name: str, *, input: Any = None, output: Any = None, error: str | None = None,
            kind: str = "action", duration_ms: float | None = None, metadata: dict | None = None) -> dict:
        """Record an action that already happened (single call)."""
        return self._http.post("/v1/events", json=self._event(
            name, input, kind, metadata, output=output, error=error, duration_ms=duration_ms)).json()

    # ---- async API
    async def acheck(self, name: str, input: Any = None, *, kind: str = "tool_call", metadata: dict | None = None,
                     snapshot: Callable[[Any], Any] | None = None) -> str | None:
        d = await self._adecided(name, input, kind, metadata)
        eid = self._final(name, d)
        await self._akeep(name, d, input, snapshot)
        return eid

    async def aask(self, question: str) -> str:
        """ask(), for async code."""
        d = await self._adecided("ask_user", {"question": question}, "question", None)
        if d is None:
            raise GatewayUnavailable("Squidbrake is unreachable, so nobody could be asked")
        self._final("ask_user", d)
        await self.aresult(d["event_id"], output={"answer": d.get("decision_note")})
        return d.get("decision_note") or ""

    async def _adecided(self, name: str, input: Any, kind: str, metadata: dict | None) -> dict | None:
        try:
            resp = await self._ahttp().post("/v1/events", json=self._event(name, input, kind, metadata))
        except httpx.TransportError:
            resp = None
        d = self._first(name, resp, input, kind)
        while d is not None and d["decision"] == "review":
            try:
                r = await self._ahttp().get(f"/v1/events/{d['event_id']}/decision",
                                          params={"wait": LONG_POLL}, timeout=LONG_POLL + 10)
                r.raise_for_status()
                d = r.json()
            except httpx.TransportError:
                if self._gave_up(d):
                    raise GatewayUnavailable(f"lost the gateway while '{name}' awaited approval")
                await asyncio.sleep(2)
        return d

    async def aresult(self, event_id: str | None, *, output: Any = None, error: str | None = None,
                      duration_ms: float | None = None) -> None:
        if event_id is None:
            return
        try:
            await self._ahttp().post(f"/v1/events/{event_id}/result",
                                   json={"output": output, "error": error, "duration_ms": duration_ms})
        except httpx.TransportError:
            pass

    # ---- decorator
    def guard(self, name: str | None = None, *, kind: str = "tool_call", snapshot: Callable[[dict], Any] | None = None):
        def deco(fn: Callable):
            tool = name or fn.__name__
            sig = inspect.signature(fn)

            def args_of(a, kw):
                bound = sig.bind_partial(*a, **kw)
                return {k: v for k, v in bound.arguments.items() if k not in ("self", "cls")}

            if inspect.iscoroutinefunction(fn):
                @functools.wraps(fn)
                async def awrapper(*a, **kw):
                    eid = await self.acheck(tool, args_of(a, kw), kind=kind, snapshot=snapshot)
                    t0 = time.perf_counter()
                    try:
                        out = await fn(*a, **kw)
                    except Exception as e:
                        await self.aresult(eid, error=f"{type(e).__name__}: {e}", duration_ms=(time.perf_counter() - t0) * 1000)
                        raise
                    await self.aresult(eid, output=out, duration_ms=(time.perf_counter() - t0) * 1000)
                    return out
                return awrapper

            @functools.wraps(fn)
            def wrapper(*a, **kw):
                eid = self.check(tool, args_of(a, kw), kind=kind, snapshot=snapshot)
                t0 = time.perf_counter()
                try:
                    out = fn(*a, **kw)
                except Exception as e:
                    self.result(eid, error=f"{type(e).__name__}: {e}", duration_ms=(time.perf_counter() - t0) * 1000)
                    raise
                self.result(eid, output=out, duration_ms=(time.perf_counter() - t0) * 1000)
                return out
            return wrapper
        return deco

    @staticmethod
    def _soft(guarded: Callable) -> Callable:
        """The guarded function, returning a refusal message instead of raising when Squidbrake stops the call."""
        if inspect.iscoroutinefunction(guarded):
            @functools.wraps(guarded)
            async def asoft(*a, **kw):
                try:
                    return await guarded(*a, **kw)
                except (Denied, GatewayUnavailable) as e:
                    return refusal(e)
            return asoft

        @functools.wraps(guarded)
        def soft(*a, **kw):
            try:
                return guarded(*a, **kw)
            except (Denied, GatewayUnavailable) as e:
                return refusal(e)
        return soft

    def guard_tools(self, tools: list, *, prefix: str = "") -> list:
        """Put every tool of an agent framework behind Squidbrake, without touching their code:

            agent = Agent(tools=gw.guard_tools([search, refund, send_email]))          # OpenAI Agents SDK
            agent = create_react_agent(llm, gw.guard_tools(tools))                      # LangChain / LangGraph

        Works on plain functions (CrewAI, PydanticAI, smolagents, Google ADK take those), LangChain tools (their
        `func` / `coroutine`) and OpenAI Agents SDK FunctionTools (their `on_invoke_tool`). The tools come back in the
        same order, as the same kind of object. A call Squidbrake stops doesn't run and returns a "NOT RUN: ..." message
        as the tool's result, so the model reads why (some frameworks end the run on a tool exception). `prefix` names
        them in the dashboard and rules, e.g. "support." """
        out = []
        for t in tools:
            name = prefix + str(getattr(t, "name", None) or getattr(t, "__name__", None) or type(t).__name__)
            if callable(getattr(t, "on_invoke_tool", None)):              # OpenAI Agents SDK FunctionTool
                original = t.on_invoke_tool

                async def invoke(ctx, raw: str, _orig=original, _name=name):
                    import json
                    try:
                        args = json.loads(raw) if raw else {}
                    except ValueError:
                        args = {"input": raw}
                    try:
                        eid = await self.acheck(_name, args)
                    except (Denied, GatewayUnavailable) as e:
                        return refusal(e)
                    t0 = time.perf_counter()
                    try:
                        res = await _orig(ctx, raw)
                    except Exception as e:
                        await self.aresult(eid, error=f"{type(e).__name__}: {e}", duration_ms=(time.perf_counter() - t0) * 1000)
                        raise
                    await self.aresult(eid, output=res, duration_ms=(time.perf_counter() - t0) * 1000)
                    return res
                t.on_invoke_tool = invoke
            elif callable(getattr(t, "func", None)) or callable(getattr(t, "coroutine", None)):   # LangChain tools
                if callable(getattr(t, "func", None)):
                    t.func = self._soft(self.guard(name)(t.func))
                if callable(getattr(t, "coroutine", None)):
                    t.coroutine = self._soft(self.guard(name)(t.coroutine))
            elif callable(t):
                t = self._soft(self.guard(name)(t))
            else:
                raise TypeError(f"don't know how to guard {t!r}: wrap its function with @gw.guard() instead")
            out.append(t)
        return out