"""
Burrow SDK - Prompt injection firewall for AI agents.

Usage:
    from burrow import BurrowGuard

    guard = BurrowGuard(
        client_id="...",
        client_secret="...",
    )
    result = guard.scan("some text to check")
    if result.is_blocked:
        print(f"Blocked: {result.category}")
"""

from __future__ import annotations

import os
import time
import uuid
from dataclasses import dataclass
from enum import Enum
from typing import Any

import httpx

from burrow._exceptions import (
    BurrowAuthError,
    BurrowBlockedError,
    BurrowError,
    BurrowQuotaError,
    BurrowTimeoutError,
)
from burrow._logging import logger
from burrow._version import __version__

__all__ = [
    "BurrowGuard",
    "ScanResult",
    "ScanAction",
    "BurrowError",
    "BurrowAuthError",
    "BurrowQuotaError",
    "BurrowBlockedError",
    "BurrowTimeoutError",
    "__version__",
]


class ScanAction(str, Enum):
    ALLOW = "allow"
    WARN = "warn"
    BLOCK = "block"


@dataclass
class ScanResult:
    """Result of a Burrow scan."""

    action: str
    confidence: float
    category: str
    request_id: str
    latency_ms: float

    @property
    def is_blocked(self) -> bool:
        return self.action == ScanAction.BLOCK

    @property
    def is_warning(self) -> bool:
        return self.action == ScanAction.WARN

    @property
    def is_allowed(self) -> bool:
        return self.action == ScanAction.ALLOW


class BurrowGuard:
    """
    Core Burrow client for scanning text for prompt injection.

    Authenticates via OAuth 2.1 client_credentials flow through Warden.
    Can be used directly or through framework-specific adapters.
    """

    def __init__(
        self,
        client_id: str | None = None,
        client_secret: str | None = None,
        api_url: str = "https://api.burrow.run",
        auth_url: str | None = None,
        timeout: float = 10.0,
        fail_open: bool = True,
        session_id: str | None = None,
    ):
        self.client_id = client_id or os.environ.get("BURROW_CLIENT_ID", "")
        self.client_secret = client_secret or os.environ.get("BURROW_CLIENT_SECRET", "")
        self.api_url = (os.environ.get("BURROW_API_URL") or api_url).rstrip("/")
        # Auth URL points to Warden (Authorization Server) for OAuth token issuance.
        # Default: {api_url}/v1/auth (production gateway proxies to Warden).
        # Local dev override: BURROW_AUTH_URL=http://localhost:8002/oauth
        self.auth_url = (
            auth_url or os.environ.get("BURROW_AUTH_URL") or f"{self.api_url}/v1/auth"
        ).rstrip("/")
        self.fail_open = fail_open
        self.session_id = session_id or str(uuid.uuid4())
        self._timeout = timeout

        # Token state
        self._access_token: str | None = None
        self._token_expires_at: float = 0.0

        # HTTP clients (lazily initialized)
        self._client: httpx.Client | None = None
        self._async_client: httpx.AsyncClient | None = None

    def _obtain_token(self) -> str:
        """Obtain a JWT access token from Warden (Authorization Server).

        Raises BurrowAuthError on 401/403 — these are permanent credential
        failures and must never be swallowed by fail_open.
        """
        resp = httpx.post(
            f"{self.auth_url}/token",
            data={
                "grant_type": "client_credentials",
                "client_id": self.client_id,
                "client_secret": self.client_secret,
            },
            timeout=self._timeout,
        )
        if resp.status_code in (401, 403):
            raise BurrowAuthError()
        resp.raise_for_status()
        data = resp.json()
        token: str = data["access_token"]
        self._access_token = token
        expires_in = data.get("expires_in", 3600)
        self._token_expires_at = time.monotonic() + expires_in - 60
        logger.info("burrow: obtained access token (expires in %ds)", expires_in)
        return token

    async def _obtain_token_async(self) -> str:
        """Obtain a JWT access token from Warden (Authorization Server, async).

        Raises BurrowAuthError on 401/403 — these are permanent credential
        failures and must never be swallowed by fail_open.
        """
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(
                f"{self.auth_url}/token",
                data={
                    "grant_type": "client_credentials",
                    "client_id": self.client_id,
                    "client_secret": self.client_secret,
                },
            )
            if resp.status_code in (401, 403):
                raise BurrowAuthError()
            resp.raise_for_status()
            data = resp.json()
        token: str = data["access_token"]
        self._access_token = token
        expires_in = data.get("expires_in", 3600)
        self._token_expires_at = time.monotonic() + expires_in - 60
        logger.info("burrow: obtained access token (expires in %ds)", expires_in)
        return token

    def _ensure_token(self) -> str:
        """Return cached access token, or obtain a new one if expired."""
        if self._access_token and time.monotonic() < self._token_expires_at:
            return self._access_token
        return self._obtain_token()

    async def _ensure_token_async(self) -> str:
        """Return cached access token, or obtain a new one if expired (async)."""
        if self._access_token and time.monotonic() < self._token_expires_at:
            return self._access_token
        return await self._obtain_token_async()

    def _get_auth_headers(self) -> dict[str, str]:
        """Get Authorization header with current JWT."""
        token = self._ensure_token()
        return {"Authorization": f"Bearer {token}"}

    def _get_client(self) -> httpx.Client:
        """Get or create sync HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.Client(
                base_url=self.api_url,
                timeout=self._timeout,
                transport=httpx.HTTPTransport(retries=2),
            )
        return self._client

    def _get_async_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._async_client is None or self._async_client.is_closed:
            self._async_client = httpx.AsyncClient(
                base_url=self.api_url,
                timeout=self._timeout,
                transport=httpx.AsyncHTTPTransport(retries=2),
            )
        return self._async_client

    @staticmethod
    def _check_quota_error(resp: httpx.Response) -> None:
        """Raise BurrowQuotaError if the response indicates quota exhaustion.

        Called before raise_for_status() so we can inspect the body and provide
        a specific error type. Non-quota 403s fall through unchanged.
        """
        if resp.status_code == 403:
            try:
                body = resp.json()
            except Exception:
                body = {}
            if body.get("error") == "quota_exceeded":
                msg = body.get("message")
                raise BurrowQuotaError(msg) if isinstance(msg, str) and msg else BurrowQuotaError()

    def scan(
        self,
        text: str,
        content_type: str = "user_prompt",
        agent: str | None = None,
        tool_name: str | None = None,
        tool_args: dict[str, Any] | None = None,
        resource: str | None = None,
    ) -> ScanResult:
        """
        Scan text for prompt injection. Synchronous.

        Args:
            text: The text to scan.
            content_type: One of "user_prompt", "tool_call", "tool_response".
            agent: Optional agent identifier for context.
            tool_name: Optional tool name for context.
            tool_args: Optional dict of tool parameters (e.g. command, file_path).
            resource: Optional resource identifier (file path, URL, command).

        Returns:
            ScanResult with action, confidence, and category.
        """
        try:
            client = self._get_client()
            headers = self._get_auth_headers()
            payload = {
                "text": text,
                "context": {
                    "type": content_type,
                    "agent": agent,
                    "tool_name": tool_name,
                    "session_id": self.session_id,
                    "tool_args": tool_args,
                    "resource": resource,
                },
            }
            resp = client.post("/v1/scan", json=payload, headers=headers)
            self._check_quota_error(resp)
            resp.raise_for_status()
            data = resp.json()
            result = ScanResult(
                action=data["action"],
                confidence=data["confidence"],
                category=data["category"],
                request_id=data["request_id"],
                latency_ms=data["latency_ms"],
            )
            logger.debug(
                "burrow scan: action=%s confidence=%.2f category=%s latency=%.1fms",
                result.action,
                result.confidence,
                result.category,
                result.latency_ms,
            )
            return result
        except BurrowAuthError:
            raise  # credential failures always propagate, never swallowed by fail_open
        except BurrowQuotaError:
            raise  # quota exhaustion always propagates, never swallowed by fail_open
        except (
            httpx.HTTPStatusError,
            httpx.ConnectError,
            httpx.TimeoutException,
            httpx.RequestError,
        ) as exc:
            logger.warning("burrow scan failed: %s", exc)
            if self.fail_open:
                return ScanResult(
                    action="allow",
                    confidence=0.0,
                    category="error",
                    request_id="",
                    latency_ms=0.0,
                )
            raise BurrowError(f"Scan request failed: {exc}") from exc

    async def scan_async(
        self,
        text: str,
        content_type: str = "user_prompt",
        agent: str | None = None,
        tool_name: str | None = None,
        tool_args: dict[str, Any] | None = None,
        resource: str | None = None,
    ) -> ScanResult:
        """
        Scan text for prompt injection. Async.

        Args:
            text: The text to scan.
            content_type: One of "user_prompt", "tool_call", "tool_response".
            agent: Optional agent identifier for context.
            tool_name: Optional tool name for context.
            tool_args: Optional dict of tool parameters (e.g. command, file_path).
            resource: Optional resource identifier (file path, URL, command).

        Returns:
            ScanResult with action, confidence, and category.
        """
        try:
            client = self._get_async_client()
            token = await self._ensure_token_async()
            headers = {"Authorization": f"Bearer {token}"}
            payload = {
                "text": text,
                "context": {
                    "type": content_type,
                    "agent": agent,
                    "tool_name": tool_name,
                    "session_id": self.session_id,
                    "tool_args": tool_args,
                    "resource": resource,
                },
            }
            resp = await client.post("/v1/scan", json=payload, headers=headers)
            self._check_quota_error(resp)
            resp.raise_for_status()
            data = resp.json()
            result = ScanResult(
                action=data["action"],
                confidence=data["confidence"],
                category=data["category"],
                request_id=data["request_id"],
                latency_ms=data["latency_ms"],
            )
            logger.debug(
                "burrow scan: action=%s confidence=%.2f category=%s latency=%.1fms",
                result.action,
                result.confidence,
                result.category,
                result.latency_ms,
            )
            return result
        except BurrowAuthError:
            raise  # credential failures always propagate, never swallowed by fail_open
        except BurrowQuotaError:
            raise  # quota exhaustion always propagates, never swallowed by fail_open
        except (
            httpx.HTTPStatusError,
            httpx.ConnectError,
            httpx.TimeoutException,
            httpx.RequestError,
        ) as exc:
            logger.warning("burrow scan failed: %s", exc)
            if self.fail_open:
                return ScanResult(
                    action="allow",
                    confidence=0.0,
                    category="error",
                    request_id="",
                    latency_ms=0.0,
                )
            raise BurrowError(f"Scan request failed: {exc}") from exc

    def close(self):
        """Close HTTP clients."""
        if self._client and not self._client.is_closed:
            self._client.close()

    async def aclose(self):
        """Close async HTTP client."""
        if self._client and not self._client.is_closed:
            self._client.close()
        if self._async_client and not self._async_client.is_closed:
            await self._async_client.aclose()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.aclose()
