"""
psiguard/client.py

The PsiGuard Python SDK client.
Wraps the PsiGuard API so developers can monitor AI responses
in their own applications with a single function call.

Usage:
    from psiguard import PsiGuard

    pg = PsiGuard(api_key="pg_...")

    result = pg.check(
        prompt="What is the capital of France?",
        response="The capital of France is Paris.",
        thread_id="session-abc123",
    )

    if result.should_block:
        print("Blocked:", result.explanation)
    else:
        print("Safe — risk score:", result.risk_score)
"""

import os
from dataclasses import dataclass, field
from typing import List, Optional

import requests

DEFAULT_BASE_URL = "https://bumpr-ai-safety-rafj.onrender.com"
DEFAULT_TIMEOUT  = 15  # seconds


# ── Response dataclasses ──────────────────────────────────────────────────────

@dataclass
class CheckResult:
    """
    Returned by PsiGuard.check().

    Attributes:
        should_block  True if PsiGuard recommends blocking this response.
        state         One of: WARMING_UP, STABLE, ANALYZING, UNSTABLE,
                      ANOMALY, DRIFT, CRITICAL, BLOCKED.
        risk_score    Integer 0-100. Higher = more risk.
        confidence    Float 0.0-1.0. How confident the engine is in its verdict.
        warnings      List of human-readable warning strings.
        explanation   Plain-English explanation of the decision.
        timestamp     Unix timestamp from the server.
    """
    should_block: bool
    state:        str
    risk_score:   int
    confidence:   float
    warnings:     List[str]
    explanation:  str
    timestamp:    float


@dataclass
class ThreadSummary:
    """
    Returned by PsiGuard.thread_summary().

    Attributes:
        total_steps  Number of checks run in this thread.
        mean_risk    Average risk score across all steps (0-100).
        max_risk     Highest risk score seen in this thread (0-100).
        state_counts Dict of state label -> count (may be empty on some plans).
    """
    total_steps:  int
    mean_risk:    float
    max_risk:     int
    state_counts: dict = field(default_factory=dict)


# ── Exceptions ────────────────────────────────────────────────────────────────

class PsiGuardError(Exception):
    """Raised when the PsiGuard API returns an error."""
    def __init__(self, message: str, status_code: Optional[int] = None):
        super().__init__(message)
        self.status_code = status_code


class PsiGuardAuthError(PsiGuardError):
    """Raised when the API key is missing, invalid, or inactive."""
    pass


class PsiGuardRateLimitError(PsiGuardError):
    """Raised when the monthly check limit is reached (Playground plan)."""
    pass


# ── Main client ───────────────────────────────────────────────────────────────

class PsiGuard:
    """
    PsiGuard AI Safety Monitoring client.

    Args:
        api_key   Your pg_... API key. If omitted, reads PSIGUARD_API_KEY
                  from the environment.
        base_url  Override the API endpoint (useful for testing).
        timeout   HTTP request timeout in seconds (default: 15).

    Example:
        pg = PsiGuard(api_key="pg_...")
        result = pg.check(prompt="...", response="...", thread_id="session-1")
    """

    def __init__(
        self,
        api_key:  Optional[str] = None,
        base_url: str           = DEFAULT_BASE_URL,
        timeout:  int           = DEFAULT_TIMEOUT,
    ):
        resolved_key = api_key or os.environ.get("PSIGUARD_API_KEY", "")
        if not resolved_key:
            raise PsiGuardAuthError(
                "No API key provided. Pass api_key='pg_...' or set the "
                "PSIGUARD_API_KEY environment variable."
            )
        if not resolved_key.startswith("pg_"):
            raise PsiGuardAuthError(
                "Invalid API key format. Keys must start with 'pg_'. "
                "Find your key in the PsiGuard dashboard."
            )

        self._api_key  = resolved_key
        self._base_url = base_url.rstrip("/")
        self._timeout  = timeout
        self._session  = requests.Session()
        self._session.headers.update({
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type":  "application/json",
            "User-Agent":    "psiguard-python/0.1.0",
        })

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _post(self, path: str, body: dict) -> dict:
        url = f"{self._base_url}{path}"
        try:
            resp = self._session.post(url, json=body, timeout=self._timeout)
        except requests.exceptions.ConnectionError as e:
            raise PsiGuardError(f"Could not connect to PsiGuard API: {e}")
        except requests.exceptions.Timeout:
            raise PsiGuardError(
                f"Request to PsiGuard API timed out after {self._timeout}s."
            )

        if resp.status_code == 401:
            raise PsiGuardAuthError(
                "Invalid or inactive API key. Check your key in the PsiGuard dashboard.",
                status_code=401,
            )
        if resp.status_code == 429:
            data = resp.json() if resp.content else {}
            raise PsiGuardRateLimitError(
                data.get("error", "Monthly limit reached. Upgrade to Pro for unlimited monitoring."),
                status_code=429,
            )
        if resp.status_code == 503:
            raise PsiGuardError(
                "PsiGuard monitoring engine is temporarily unavailable. Try again shortly.",
                status_code=503,
            )
        if not resp.ok:
            try:
                msg = resp.json().get("error", resp.text)
            except Exception:
                msg = resp.text
            raise PsiGuardError(f"API error ({resp.status_code}): {msg}", status_code=resp.status_code)

        return resp.json()

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

    def check(
        self,
        prompt:         str,
        response:       str,
        thread_id:      str = "default",
        risk_threshold: int = 70,
    ) -> CheckResult:
        """
        Monitor an AI response and get a safety decision.

        Args:
            prompt          The user's input / question sent to the AI.
            response        The AI's response to evaluate.
            thread_id       Identifies the conversation session. Use the same
                            thread_id across turns to enable drift and coherence
                            tracking.
            risk_threshold  0-100. Responses with risk_score above this will
                            have should_block=True. Default: 70.

        Returns:
            CheckResult

        Raises:
            PsiGuardAuthError       Bad or missing API key.
            PsiGuardRateLimitError  Monthly limit reached (Playground plan).
            PsiGuardError           Any other API or network error.

        Example:
            result = pg.check(
                prompt="Explain quantum tunneling.",
                response="Quantum tunneling allows particles to...",
                thread_id="chat-session-42",
            )
            if result.should_block:
                raise ValueError("Unsafe AI response detected")
        """
        if not prompt or not prompt.strip():
            raise ValueError("prompt must be a non-empty string.")
        if not response or not response.strip():
            raise ValueError("response must be a non-empty string.")

        data = self._post("/v1/check", {
            "prompt":         prompt.strip(),
            "response":       response.strip(),
            "thread_id":      thread_id or "default",
            "risk_threshold": int(risk_threshold),
        })

        return CheckResult(
            should_block = bool(data.get("should_block", False)),
            state        = str(data.get("state", "UNKNOWN")),
            risk_score   = int(data.get("risk_score", 0)),
            confidence   = float(data.get("confidence", 0.0)),
            warnings     = list(data.get("warnings", [])),
            explanation  = str(data.get("explanation", "")),
            timestamp    = float(data.get("timestamp", 0.0)),
        )

    def thread_reset(self, thread_id: str) -> bool:
        """
        Clear a conversation thread's monitoring history.

        Call this when a new conversation starts so PsiGuard doesn't carry
        over context from a previous session using the same thread_id.

        Args:
            thread_id  Must match a thread_id previously used in check().

        Returns:
            True if the thread was found and reset.

        Raises:
            PsiGuardError  Thread not found or network error.
        """
        if not thread_id or not thread_id.strip():
            raise ValueError("thread_id must be a non-empty string.")
        data = self._post("/v1/thread/reset", {"thread_id": thread_id.strip()})
        return bool(data.get("ok", False))

    def thread_summary(self, thread_id: str) -> ThreadSummary:
        """
        Get aggregated monitoring stats for a conversation thread.

        Args:
            thread_id  The thread to summarise.

        Returns:
            ThreadSummary with total_steps, mean_risk, max_risk, state_counts.

        Raises:
            PsiGuardError  Thread not found or network error.
        """
        if not thread_id or not thread_id.strip():
            raise ValueError("thread_id must be a non-empty string.")
        data = self._post("/v1/thread/summary", {"thread_id": thread_id.strip()})
        return ThreadSummary(
            total_steps  = int(data.get("total_steps", 0)),
            mean_risk    = float(data.get("mean_risk", 0.0)),
            max_risk     = int(data.get("max_risk", 0)),
            state_counts = dict(data.get("state_counts", {})),
        )
