#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PsiGuard Python SDK
===================

The official Python client for the PsiGuard guarded-chat API.

PsiGuard wraps a language model with a live structural monitor. You send it a
prompt; it runs the model, watches the generation as it happens, and hands you
back a clean answer plus one coarse signal telling you whether it had to step
in. Everything about *how* it decides stays on PsiGuard's servers — this client
only ever sees the answer and that one signal.

What this client does
----------------------
  * Authenticates with your PsiGuard API key.
  * Sends a prompt to the guard endpoint and returns the guarded answer.
  * Keeps a conversation together across turns (so monitoring has context).
  * Turns HTTP errors into clear, catchable Python exceptions.
  * Retries the handful of failures that are safe to retry, and never the
    ones that aren't.

What this client is NOT
-----------------------
  It is a courier, not the engine. It contains none of PsiGuard's inner
  workings — no metrics, no thresholds, no scoring. It speaks only the public
  contract: a prompt goes out, an answer and a protection signal come back.

Quick start
-----------
    pip install requests        # the only dependency

    from psiguard import Client

    client = Client(api_key="psg_live_your_key_here")
    result = client.guard("What is your return policy?")
    print(result.answer)        # the reply to show your user
    print(result.protection)    # "safe" | "caution" | "intervened"

Set PSIGUARD_API_KEY in your environment and you can skip passing the key:

    client = Client()           # reads PSIGUARD_API_KEY

Multi-turn
----------
    chat = client.conversation()           # one ongoing thread
    chat.send("Hi, I need help with an order.")
    chat.send("It hasn't shipped yet.")    # remembers the turn before it

This module has no PsiGuard-specific dependencies and is safe to vendor: drop
psiguard.py into your project, ensure `requests` is installed, and import it.
"""

from __future__ import annotations

__version__ = "1.1.0"
__all__ = [
    "Client",
    "Conversation",
    "GuardResult",
    "Account",
    "Protection",
    # Terminal display (optional; the client itself never prints)
    "supports_color",
    "badge",
    "format_result",
    # Exceptions
    "PsiGuardError",
    "ConfigurationError",
    "APIConnectionError",
    "APIStatusError",
    "BadRequestError",
    "AuthenticationError",
    "RateLimitError",
    "UsageLimitError",
    "ServerError",
]

import os
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

try:
    import requests
except ImportError:  # pragma: no cover
    raise ImportError(
        "The PsiGuard SDK needs the 'requests' package.\n"
        "Install it with:  pip install requests"
    )


# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

DEFAULT_BASE_URL = "https://psiguard.net"
DEFAULT_TIMEOUT = 60.0          # seconds; a guarded turn runs a full model call
DEFAULT_MAX_RETRIES = 2         # extra attempts for retryable failures only
_RETRY_BACKOFF = 0.75           # seconds; grows each retry (0.75, 1.5, 3.0, …)
_USER_AGENT = f"psiguard-python/{__version__}"


# ---------------------------------------------------------------------------
# The protection signal
#
# Every guarded turn comes back with exactly one of these. It says how much
# PsiGuard had to do — never anything about how it decided.
# ---------------------------------------------------------------------------

class Protection:
    """The three values the `protection` field can take.

    Use these constants instead of bare strings so a typo is a NameError you
    catch immediately rather than a comparison that silently never matches.
    """

    SAFE = "safe"               # clean run; PsiGuard didn't need to act
    CAUTION = "caution"         # steered back on course mid-generation
    INTERVENED = "intervened"   # response withheld; answer is a neutral refusal

    ALL = (SAFE, CAUTION, INTERVENED)


# ---------------------------------------------------------------------------
# Results
# ---------------------------------------------------------------------------

@dataclass
class GuardResult:
    """The outcome of one guarded turn.

    Attributes:
        answer:          The reply to show your user. On an intervention this
                         is a neutral refusal, indistinguishable from an
                         ordinary model decline — safe to display as-is.
        protection:      One of Protection.SAFE / .CAUTION / .INTERVENED.
        conversation_id: An echo of the id this turn belongs to.
        raw:             The exact JSON object the API returned, in case a
                         future field appears before this class is updated.
    """

    answer: str
    protection: str
    conversation_id: str
    raw: Dict[str, Any] = field(default_factory=dict, repr=False)

    @property
    def safe(self) -> bool:
        """True when PsiGuard did not need to act on this turn."""
        return self.protection == Protection.SAFE

    @property
    def caution(self) -> bool:
        """True when PsiGuard steered the response back on course."""
        return self.protection == Protection.CAUTION

    @property
    def intervened(self) -> bool:
        """True when PsiGuard withheld the response (answer is a refusal)."""
        return self.protection == Protection.INTERVENED


@dataclass
class Account:
    """Who an API key belongs to, as reported by a verify check."""

    account: str
    auth_via: str
    raw: Dict[str, Any] = field(default_factory=dict, repr=False)


# ---------------------------------------------------------------------------
# Exceptions
#
# One base class so callers can `except PsiGuardError` and be done, plus
# specific subclasses for the cases worth handling differently.
# ---------------------------------------------------------------------------

class PsiGuardError(Exception):
    """Base class for everything this SDK raises."""


class ConfigurationError(PsiGuardError):
    """The client was set up wrong (e.g. no API key). No request was made."""


class APIConnectionError(PsiGuardError):
    """The request never got a reply — network error, DNS, or timeout.

    These are safe to retry; the SDK already retried up to `max_retries`
    before raising.
    """


class APIStatusError(PsiGuardError):
    """The API replied with a non-success HTTP status.

    Attributes:
        status_code: The HTTP status (e.g. 400, 401, 429, 502).
        error:       The short machine-readable code from the body's `error`
                     field, when present (e.g. "missing 'prompt'").
        message:     A human-readable message from the body, when present.
        body:        The full parsed response body.
    """

    def __init__(
        self,
        message: str,
        *,
        status_code: int,
        error: Optional[str] = None,
        body: Optional[Dict[str, Any]] = None,
    ) -> None:
        super().__init__(message)
        self.status_code = status_code
        self.error = error
        self.message = message
        self.body = body or {}


class BadRequestError(APIStatusError):
    """400 — the request was malformed or the account has no model configured."""


class AuthenticationError(APIStatusError):
    """401 — the API key was missing, invalid, or revoked."""


class RateLimitError(APIStatusError):
    """429 — too many requests in a short window. Slow down and retry shortly."""


class UsageLimitError(APIStatusError):
    """429 — a usage limit was reached (a stop sign, not a retry).

    This is distinct from RateLimitError: you've hit either the account's
    monthly cap or a provider budget. Retrying won't help — raise the cap or
    bring your own provider key.

    Attributes (any may be None depending on which limit was hit):
        scope:    "account" for a monthly cap, "provider" for a provider budget.
        plan:     The account's plan (account cap only).
        provider: The provider whose budget was reached (provider budget only).
        used:     How much has been used this month.
        limit:    The cap or budget that was reached.
        month:    The month the counter applies to.
    """

    def __init__(self, message: str, **kwargs: Any) -> None:
        body = kwargs.get("body") or {}
        super().__init__(message, **kwargs)
        code = body.get("error")
        self.scope = "provider" if code == "provider_budget_reached" else "account"
        self.plan = body.get("plan")
        self.provider = body.get("provider")
        self.used = body.get("used")
        # The body names the ceiling "cap" for accounts and "budget" for
        # providers; expose both under one friendly name.
        self.limit = body.get("cap", body.get("budget"))
        self.month = body.get("month")


class ServerError(APIStatusError):
    """5xx — the server or the underlying model failed. Usually safe to retry."""


# The set of usage-cap error codes that mean "stop", not "slow down".
_USAGE_CODES = {"monthly_cap_reached", "provider_budget_reached"}


# ---------------------------------------------------------------------------
# The client
# ---------------------------------------------------------------------------

class Client:
    """A connection to the PsiGuard API.

    Args:
        api_key:     Your PsiGuard API key (looks like "psg_live_…"). Falls
                     back to the PSIGUARD_API_KEY environment variable.
        base_url:    Your PsiGuard host. Falls back to the PSIGUARD_BASE_URL
                     environment variable, then to https://psiguard.net.
        timeout:     Per-request timeout in seconds. A guarded turn runs a
                     full model call, so the default is generous.
        max_retries: How many extra times to retry a *retryable* failure
                     (connection errors and 5xx). 4xx and usage caps are
                     never retried.
        session:     An optional requests.Session to reuse. One is created if
                     you don't pass your own; it's closed by .close() / the
                     context manager only when the client made it.

    The client is safe to keep and reuse for the life of your process; it
    keeps a connection pool open under the hood. Use it as a context manager
    to close that pool when you're done:

        with Client() as client:
            client.guard("hello")
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        *,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
        session: Optional["requests.Session"] = None,
    ) -> None:
        key = api_key or os.environ.get("PSIGUARD_API_KEY")
        if not key:
            raise ConfigurationError(
                "No API key. Pass api_key=... or set the PSIGUARD_API_KEY "
                "environment variable. Create a key at "
                "https://psiguard.net/keys"
            )
        self.api_key = key.strip()

        host = base_url or os.environ.get("PSIGUARD_BASE_URL") or DEFAULT_BASE_URL
        self.base_url = host.rstrip("/")

        self.timeout = float(timeout)
        self.max_retries = max(0, int(max_retries))

        # Track whether we own the session so we don't close one the caller
        # handed us.
        self._owns_session = session is None
        self._session = session or requests.Session()

    # -- lifecycle ----------------------------------------------------------

    def close(self) -> None:
        """Close the underlying connection pool (only if the client made it)."""
        if self._owns_session:
            self._session.close()

    def __enter__(self) -> "Client":
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()

    # -- public API ---------------------------------------------------------

    def guard(
        self,
        prompt: str,
        *,
        system_prompt: Optional[str] = None,
        conversation_id: Optional[str] = None,
        history: Optional[List[Dict[str, str]]] = None,
    ) -> GuardResult:
        """Run one guarded turn and return the answer.

        This is the call most integrations only ever need.

        Args:
            prompt:          The user's message for this turn. Required.
            system_prompt:   Your assistant's personality, scope, or
                             instructions. Omit to run the model with no
                             steering (still guarded).
            conversation_id: A caller-chosen id that groups turns into one
                             thread, so monitoring carries context across a
                             conversation. Reuse the same id for every turn in
                             a chat. Omit and the server uses "default".
            history:         Prior turns you want the model to see, as
                             {"role": ..., "content": ...} objects. You own and
                             store the transcript; PsiGuard doesn't retain it.

        Returns:
            A GuardResult with .answer, .protection, and .conversation_id.

        Raises:
            BadRequestError, AuthenticationError, RateLimitError,
            UsageLimitError, ServerError, APIConnectionError — see each.
        """
        if not isinstance(prompt, str) or not prompt.strip():
            raise ConfigurationError("guard() needs a non-empty 'prompt' string.")

        body: Dict[str, Any] = {"prompt": prompt}
        if system_prompt:
            body["system_prompt"] = system_prompt
        if conversation_id:
            body["conversation_id"] = conversation_id
        if history:
            body["history"] = self._clean_history(history)

        data = self._request("POST", "/v1/guard", json_body=body)

        protection = data.get("protection", Protection.SAFE)
        return GuardResult(
            answer=data.get("answer", ""),
            protection=protection,
            conversation_id=data.get("conversation_id", conversation_id or "default"),
            raw=data,
        )

    def whoami(self) -> Account:
        """Confirm the API key is valid and see which account it belongs to.

        Lightweight: runs no model and isn't metered. Handy as a setup check
        or health probe. Raises AuthenticationError if the key is bad.
        """
        data = self._request("GET", "/v1/whoami")
        return Account(
            account=data.get("account", ""),
            auth_via=data.get("auth_via", ""),
            raw=data,
        )

    def conversation(
        self,
        conversation_id: Optional[str] = None,
        *,
        system_prompt: Optional[str] = None,
    ) -> "Conversation":
        """Start a multi-turn conversation that manages its own thread.

        The returned Conversation reuses one conversation_id and accumulates
        history for you, so each .send() carries the turns before it.

        Args:
            conversation_id: The thread id to use. A random one is generated
                             if you don't supply it.
            system_prompt:   An assistant personality applied to every turn in
                             this conversation.
        """
        return Conversation(
            client=self,
            conversation_id=conversation_id or _new_conversation_id(),
            system_prompt=system_prompt,
        )

    # -- internals ----------------------------------------------------------

    @staticmethod
    def _clean_history(history: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """Keep only well-formed {role, content} string pairs.

        Mirrors what the server accepts, so malformed entries fail fast here
        with a clear message instead of being silently dropped over the wire.
        """
        cleaned: List[Dict[str, str]] = []
        for i, m in enumerate(history):
            if (
                isinstance(m, dict)
                and isinstance(m.get("role"), str)
                and m.get("role")
                and isinstance(m.get("content"), str)
                and m.get("content")
            ):
                cleaned.append({"role": m["role"], "content": m["content"]})
            else:
                raise ConfigurationError(
                    f"history[{i}] must be a dict with non-empty string "
                    f"'role' and 'content'."
                )
        return cleaned

    @property
    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
            "User-Agent": _USER_AGENT,
        }

    def _request(
        self,
        method: str,
        path: str,
        *,
        json_body: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Make one API call, with retries on the failures that are safe to retry.

        Retried: connection errors, timeouts, and 5xx responses.
        Never retried: 4xx (including auth and usage caps) and plain rate
        limits — retrying those just wastes a call.
        """
        url = self.base_url + path
        attempt = 0

        while True:
            attempt += 1
            try:
                resp = self._session.request(
                    method,
                    url,
                    headers=self._headers,
                    json=json_body,
                    timeout=self.timeout,
                )
            except requests.RequestException as e:
                # No response at all. Retry if we have attempts left.
                if attempt <= self.max_retries:
                    _sleep_backoff(attempt)
                    continue
                raise APIConnectionError(
                    f"Could not reach PsiGuard at {url}: {e}"
                ) from e

            # A 5xx is the model or server faltering — retry it.
            if resp.status_code >= 500 and attempt <= self.max_retries:
                _sleep_backoff(attempt)
                continue

            return self._handle_response(resp)

    @staticmethod
    def _handle_response(resp: "requests.Response") -> Dict[str, Any]:
        """Turn an HTTP response into a dict, or raise the right exception."""
        # Parse the body defensively; some errors may not carry JSON.
        try:
            body = resp.json()
            if not isinstance(body, dict):
                body = {"raw": body}
        except ValueError:
            body = {}

        status = resp.status_code

        if 200 <= status < 300:
            return body

        code = body.get("error")
        message = (
            body.get("message")
            or (code if isinstance(code, str) else None)
            or f"PsiGuard request failed with status {status}"
        )
        kwargs = dict(status_code=status, error=code, body=body)

        if status == 400:
            raise BadRequestError(message, **kwargs)
        if status == 401:
            raise AuthenticationError(message, **kwargs)
        if status == 429:
            if code in _USAGE_CODES:
                raise UsageLimitError(message, **kwargs)
            raise RateLimitError(message, **kwargs)
        if status >= 500:
            raise ServerError(message, **kwargs)

        # Any other non-success status: surface it plainly.
        raise APIStatusError(message, **kwargs)


# ---------------------------------------------------------------------------
# Terminal display helpers (optional)
#
# The client itself prints NOTHING and never puts color codes inside
# GuardResult.answer. That string is your user's answer: it may be written to
# a log, stored in a database, or rendered in a browser, and ANSI escape
# sequences in any of those places show up as garbage like "[32m". So the
# color lives out here, in helpers you opt into when you are looking at a
# terminal — and nowhere else.
#
#     from psiguard import Client, format_result
#     result = Client().guard("hello")
#     print(format_result(result))
#
# Color is applied only when the output really is an interactive terminal.
# Piping to a file, redirecting into a log, or running in CI all fall back to
# plain text automatically, so nothing downstream ever eats an escape code.
# NO_COLOR=1 disables it; PSIGUARD_COLOR=1 forces it on.
# ---------------------------------------------------------------------------

import sys

_RESET = "\033[0m"
_BOLD = "\033[1m"
_DIM = "\033[2m"

# One color per protection value, plus the glyph and plain-language caption
# that ride along with it. Green / amber / red, read at a glance.
_STYLE: Dict[str, Dict[str, str]] = {
    Protection.SAFE: {
        "color": "\033[38;5;42m",    # green
        "glyph": "\u25cf",            # ●
        "label": "SAFE",
        "caption": "clean run \u2014 PsiGuard didn't need to act",
    },
    Protection.CAUTION: {
        "color": "\033[38;5;214m",   # amber
        "glyph": "\u25b2",            # ▲
        "label": "CAUTION",
        "caption": "steered back on course mid-answer",
    },
    Protection.INTERVENED: {
        "color": "\033[38;5;203m",   # red
        "glyph": "\u25a0",            # ■
        "label": "INTERVENED",
        "caption": "response withheld \u2014 the answer below is a refusal",
    },
}


def supports_color(stream: Any = None) -> bool:
    """True when it's safe to write ANSI color to `stream` (default stdout).

    False for pipes, files, CI logs, and dumb terminals — which is exactly
    when color would turn into visible junk instead of visible meaning.
    """
    stream = stream if stream is not None else sys.stdout

    forced = os.environ.get("PSIGUARD_COLOR", "").strip().lower()
    if forced in ("1", "true", "yes", "always"):
        return True
    if forced in ("0", "false", "no", "never"):
        return False
    if os.environ.get("NO_COLOR") is not None:
        return False
    if os.environ.get("TERM", "").lower() == "dumb":
        return False

    try:
        if not stream.isatty():
            return False
    except Exception:
        return False

    # Old Windows consoles don't understand ANSI. Modern Terminal, VS Code,
    # and anything running under ANSICON do.
    if os.name == "nt":
        return bool(
            os.environ.get("WT_SESSION")
            or os.environ.get("ANSICON")
            or os.environ.get("TERM_PROGRAM") == "vscode"
        )
    return True


def badge(protection: str, *, color: Optional[bool] = None) -> str:
    """A one-line status chip: '● SAFE'.

    Args:
        protection: A GuardResult.protection value.
        color:      Force color on/off. Default: decide from the terminal.
    """
    style = _STYLE.get(protection)
    if style is None:  # a value this SDK version doesn't know yet
        return str(protection).upper()

    text = f"{style['glyph']} {style['label']}"
    use_color = supports_color() if color is None else bool(color)
    return f"{style['color']}{_BOLD}{text}{_RESET}" if use_color else text


def format_result(
    result: GuardResult,
    *,
    color: Optional[bool] = None,
    gutter: bool = True,
) -> str:
    """Render a guarded turn for a terminal: status chip, then the answer.

    The answer sits behind a colored gutter rail, so the status is impossible
    to miss without tinting a single character of the text your user reads.

        ● SAFE  clean run — PsiGuard didn't need to act
        │ Our return policy allows returns within 30 days...

    Args:
        result: The GuardResult from client.guard() or chat.send().
        color:  Force color on/off. Default: decide from the terminal.
        gutter: Draw the vertical rail beside the answer. Off gives you a
                bare chip plus the plain answer.
    """
    style = _STYLE.get(result.protection, _STYLE[Protection.SAFE])
    use_color = supports_color() if color is None else bool(color)

    chip = badge(result.protection, color=use_color)
    caption = style["caption"]
    if use_color:
        caption = f"{_DIM}{caption}{_RESET}"

    lines = [f"{chip}  {caption}"]

    answer = (result.answer or "").rstrip()
    if gutter:
        rail = "\u2502"  # │
        if use_color:
            rail = f"{style['color']}{rail}{_RESET}"
        for line in answer.splitlines() or [""]:
            lines.append(f"{rail} {line}")
    else:
        lines.append(answer)

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Multi-turn helper
# ---------------------------------------------------------------------------

class Conversation:
    """An ongoing guarded conversation on a single thread.

    Holds a conversation_id and a running transcript so each turn carries the
    ones before it. You normally make one with `client.conversation()` rather
    than constructing it directly.

        chat = client.conversation(system_prompt="You are Acme's assistant.")
        first = chat.send("Hi!")
        second = chat.send("What are your hours?")   # sees the first turn

    The transcript lives in .history as {"role", "content"} dicts. It grows by
    two entries per turn (the user prompt and the assistant's answer). On an
    intervention, the neutral refusal is what gets recorded as the assistant
    turn — the same thing your user saw.
    """

    def __init__(
        self,
        client: Client,
        conversation_id: str,
        *,
        system_prompt: Optional[str] = None,
    ) -> None:
        self.client = client
        self.conversation_id = conversation_id
        self.system_prompt = system_prompt
        self.history: List[Dict[str, str]] = []

    def send(self, prompt: str) -> GuardResult:
        """Send the next turn and record both sides in the transcript."""
        result = self.client.guard(
            prompt,
            system_prompt=self.system_prompt,
            conversation_id=self.conversation_id,
            history=self.history,
        )
        self.history.append({"role": "user", "content": prompt})
        self.history.append({"role": "assistant", "content": result.answer})
        return result

    def reset(self) -> None:
        """Forget the transcript but keep the same conversation_id."""
        self.history = []

    def __repr__(self) -> str:  # pragma: no cover
        return (
            f"Conversation(conversation_id={self.conversation_id!r}, "
            f"turns={len(self.history) // 2})"
        )


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------

def _new_conversation_id() -> str:
    """A short, unique-enough id for a fresh conversation thread."""
    return "conv-" + uuid.uuid4().hex[:16]


def _sleep_backoff(attempt: int) -> None:
    """Wait a growing amount before the next retry (0.75s, 1.5s, 3s, …)."""
    time.sleep(_RETRY_BACKOFF * (2 ** (attempt - 1)))


# ---------------------------------------------------------------------------
# A tiny smoke test you can run directly:
#     PSIGUARD_API_KEY=psg_live_... python psiguard.py "Hello there"
# ---------------------------------------------------------------------------

if __name__ == "__main__":  # pragma: no cover
    prompt = " ".join(sys.argv[1:]) or "Say hello in one short sentence."
    try:
        with Client() as _c:
            who = _c.whoami()
            print(f"Authenticated as account: {who.account}\n")
            res = _c.guard(prompt)
            print(format_result(res))
    except PsiGuardError as err:
        print(f"PsiGuard error: {err}")
        sys.exit(1)