"""
SovereignShield API Client
~~~~~~~~~~~~~~~~~~~~~~~~~~~

Core client class for interacting with the SovereignShield API.

Copyright (c) 2026 Mattijs Moens. All rights reserved.
"""

import time
from typing import Any, Dict, Optional

import requests

from sovereign_shield_client.exceptions import (
    SovereignShieldError,
    AuthenticationError,
    RateLimitError,
    QuotaExceededError,
    ValidationError,
    ServerError,
    InputBlockedError,
)
from sovereign_shield_client.models import ScanResult, UsageInfo, ReportResult

_DEFAULT_BASE_URL = "https://api.sovereign-shield.net"
_DEFAULT_TIMEOUT = 10.0
_DEFAULT_MAX_RETRIES = 3
_RETRY_BACKOFF = 1.0  # seconds, doubles each retry


class SovereignShield:
    """Official Python client for the SovereignShield API.

    Args:
        api_key: Your SovereignShield API key (starts with ``ss_``).
        base_url: API base URL. Defaults to the production endpoint.
        timeout: Request timeout in seconds. Defaults to 10.
        max_retries: Number of automatic retries on transient failures. Defaults to 3.

    Example::

        from sovereign_shield import SovereignShield

        shield = SovereignShield(api_key="ss_your_key")
        safe_input = shield.scan("user input here")
        response = your_llm.generate(safe_input)
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = _DEFAULT_BASE_URL,
        timeout: float = _DEFAULT_TIMEOUT,
        max_retries: int = _DEFAULT_MAX_RETRIES,
    ):
        if not api_key:
            raise ValueError("api_key is required")

        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._max_retries = max_retries
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "User-Agent": "sovereign-shield-python/1.0.0",
            }
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def scan(
        self,
        input: str,
        user_id: str = "anonymous",
        action: Optional[str] = None,
        payload: Optional[str] = None,
    ) -> str:
        """Scan user input through all 4 security layers.

        Returns the safe input directly as a string. If the input is
        dangerous, raises ``InputBlockedError`` instead.

        Args:
            input: Raw user input to scan.
            user_id: User identifier for firewall checks.
            action: Optional action type (e.g. ``BROWSE``, ``SHELL_EXEC``).
            payload: Optional action payload (defaults to input).

        Returns:
            The safe input string, ready to pass to your LLM.

        Raises:
            InputBlockedError: If the input is blocked by security layers.
            AuthenticationError: If the API key is invalid.
            RateLimitError: If the per-minute rate limit is exceeded.
            QuotaExceededError: If the monthly scan quota is exceeded.
        """
        body: Dict[str, Any] = {"input": input, "user_id": user_id}
        if action is not None:
            body["action"] = action
        if payload is not None:
            body["payload"] = payload

        data = self._request("POST", "/api/v1/scan", json=body)

        if not data["allowed"]:
            raise InputBlockedError(
                data["reason"],
                stage=data["stage"],
                scan_id=data["scan_id"],
            )

        return data.get("clean_input", input)

    def veto(self, response: str) -> str:
        """Scan an LLM response before it reaches the user.

        Returns the safe response directly as a string. If the response
        contains dangerous content (PII leaks, credential exposure, etc.),
        raises ``InputBlockedError`` instead.

        Args:
            response: Raw LLM response to scan.

        Returns:
            The safe response string, ready to show to the user.

        Raises:
            InputBlockedError: If the response is blocked by security layers.
            AuthenticationError: If the API key is invalid.
            RateLimitError: If the per-minute rate limit is exceeded.
            QuotaExceededError: If the monthly scan quota is exceeded.
        """
        data = self._request("POST", "/api/v1/scan/veto", json={"input": response})

        if not data["allowed"]:
            raise InputBlockedError(
                data["reason"],
                stage=data.get("layer", "VetoShield"),
                scan_id=data["scan_id"],
            )

        return response

    def usage(self) -> UsageInfo:
        """Check scan usage and account limits.

        Returns:
            UsageInfo with tier, daily/monthly counts, and remaining scans.
        """
        data = self._request("GET", "/api/v1/usage")

        return UsageInfo(
            tier=data["tier"],
            scans_today=data["scans_today"],
            scans_this_month=data["scans_this_month"],
            monthly_limit=data["monthly_limit"],
            remaining=data["remaining"],
        )

    def health(self) -> Dict[str, Any]:
        """Check API health status.

        Returns:
            Dict with status, version, shield name, and integrity flag.
        """
        return self._request("GET", "/api/v1/health")

    def report(self, scan_id: str, reason: str) -> ReportResult:
        """Report a missed attack (false negative) to improve the shield.

        Args:
            scan_id: The ``scan_id`` from the original scan response.
            reason: Why you believe this should have been blocked.

        Returns:
            ReportResult with processing status and whether a rule was auto-deployed.
        """
        data = self._request(
            "POST",
            "/api/v1/report",
            json={"scan_id": scan_id, "reason": reason},
        )

        return ReportResult(
            report_id=data["report_id"],
            status=data["status"],
            rule_created=data.get("rule_created", False),
            sandbox_result=data.get("sandbox_result"),
            message=data["message"],
        )

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _request(
        self,
        method: str,
        path: str,
        json: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Make an API request with automatic retry on transient failures."""
        url = f"{self._base_url}{path}"
        last_error = None

        for attempt in range(self._max_retries + 1):
            try:
                resp = self._session.request(
                    method,
                    url,
                    json=json,
                    timeout=self._timeout,
                )
            except requests.ConnectionError as exc:
                last_error = SovereignShieldError(
                    f"Connection failed: {exc}", status_code=0
                )
                if attempt < self._max_retries:
                    time.sleep(_RETRY_BACKOFF * (2 ** attempt))
                    continue
                raise last_error from exc
            except requests.Timeout as exc:
                last_error = SovereignShieldError(
                    f"Request timed out after {self._timeout}s", status_code=0
                )
                if attempt < self._max_retries:
                    time.sleep(_RETRY_BACKOFF * (2 ** attempt))
                    continue
                raise last_error from exc

            # Success
            if resp.status_code == 200:
                return resp.json()

            # Client errors — don't retry
            if resp.status_code == 401 or resp.status_code == 403:
                detail = self._extract_detail(resp)
                raise AuthenticationError(
                    detail or "Invalid or missing API key",
                    status_code=resp.status_code,
                )

            if resp.status_code == 422:
                detail = self._extract_detail(resp)
                raise ValidationError(
                    detail or "Invalid request payload",
                    status_code=422,
                )

            if resp.status_code == 429:
                detail = self._extract_detail(resp)
                retry_after = None
                if "Retry-After" in resp.headers:
                    try:
                        retry_after = float(resp.headers["Retry-After"])
                    except (ValueError, TypeError):
                        pass

                # Check if it's a quota error vs rate limit
                if "quota" in (detail or "").lower() or "limit" in (detail or "").lower():
                    if "month" in (detail or "").lower():
                        raise QuotaExceededError(
                            detail or "Monthly scan quota exceeded",
                            status_code=429,
                        )

                raise RateLimitError(
                    detail or "Rate limit exceeded",
                    retry_after=retry_after,
                )

            # Server errors — retry
            if resp.status_code >= 500:
                detail = self._extract_detail(resp)
                last_error = ServerError(
                    detail or f"Server error ({resp.status_code})",
                    status_code=resp.status_code,
                )
                if attempt < self._max_retries:
                    time.sleep(_RETRY_BACKOFF * (2 ** attempt))
                    continue
                raise last_error

            # Unknown error
            detail = self._extract_detail(resp)
            raise SovereignShieldError(
                detail or f"Unexpected response ({resp.status_code})",
                status_code=resp.status_code,
            )

        # Should not reach here, but just in case
        raise last_error or SovereignShieldError("Request failed after all retries")

    @staticmethod
    def _extract_detail(resp: requests.Response) -> Optional[str]:
        """Extract error detail from a JSON response."""
        try:
            body = resp.json()
            if isinstance(body, dict):
                return body.get("detail") or body.get("message") or body.get("error")
        except (ValueError, KeyError):
            pass
        return None

    def __repr__(self) -> str:
        masked = self._api_key[:6] + "..." if len(self._api_key) > 6 else "***"
        return f"SovereignShield(api_key='{masked}', base_url='{self._base_url}')"
