"""
GuardrailX Python SDK - Client

Main client for interacting with GuardrailX services.
"""

from __future__ import annotations

import json
import time
from typing import Any, Dict, Optional, Union

import requests

from .types import (
    GuardrailxClientOptions,
    GuardrailxError,
    LLMRequest,
    RehydrateRequest,
    RehydrateResponse,
    SanitizeRequest,
    SanitizeResponse,
)


class GuardrailxClient:
    """
    GuardrailX Client
    
    Provides methods to interact with GuardrailX policy-core and gateway services.
    
    Example:
        >>> client = GuardrailxClient(
        ...     api_key="gxk_your_api_key_here",
        ...     base_url="http://localhost:8080/api/v1"
        ... )
        >>> result = client.sanitize({
        ...     "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}]
        ... })
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "http://localhost:8080/api/v1",
        token: Optional[str] = None,
        tenant_id: Optional[str] = None,
        user_id: Optional[str] = None,
        timeout_ms: int = 30000,
        retries: int = 2,
        retry_delay_ms: int = 500,
        headers: Optional[Dict[str, str]] = None
    ) -> None:
        """
        Create a new GuardrailX client.
        
        Args:
            api_key: GuardrailX API key (starts with gxk_)
            base_url: Base URL for GuardrailX API (default: http://localhost:8080/api/v1)
            token: Optional JWT token for authentication
            tenant_id: Optional tenant ID for multi-tenancy
            user_id: Optional user ID for tracking
            timeout_ms: Request timeout in milliseconds (default: 30000)
            retries: Number of retry attempts (default: 2)
            retry_delay_ms: Delay between retries in milliseconds (default: 500)
            headers: Additional HTTP headers to include in requests
        """
        self._base_url = base_url.rstrip("/")
        self._token = token
        self._api_key = api_key
        self._tenant_id = tenant_id
        self._user_id = user_id
        self._timeout_ms = timeout_ms
        self._retries = retries
        self._retry_delay_ms = retry_delay_ms
        self._headers = headers or {}

    def sanitize(self, request: Union[SanitizeRequest, Dict[str, Any]]) -> SanitizeResponse:
        """
        Sanitize messages before sending to LLM (detect and tokenize sensitive data).
        
        Args:
            request: Messages to sanitize with optional metadata
            
        Returns:
            Sanitized messages with session ID for rehydration
            
        Example:
            >>> result = client.sanitize({
            ...     "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}]
            ... })
            >>> # Send result["sanitized"] to your LLM
        """
        payload = request if isinstance(request, dict) else {
            "messages": request.messages,
            "metadata": request.metadata
        }
        response = self._request("POST", "/gateway/sanitize", payload)
        return SanitizeResponse(
            sanitized=response["sanitized"],
            session_id=response["sessionId"],
            detections=response.get("detections", []),
            tokens_used=response.get("tokensUsed", 0)
        )

    def rehydrate(self, request: Union[RehydrateRequest, Dict[str, Any]]) -> RehydrateResponse:
        """
        Rehydrate LLM response (restore original sensitive data).
        
        Args:
            request: LLM response messages and session ID from sanitize
            
        Returns:
            Original messages with sensitive data restored
            
        Example:
            >>> result = client.rehydrate({
            ...     "messages": llm_response,
            ...     "sessionId": sanitize_result["sessionId"]
            ... })
            >>> # result["original"] contains restored sensitive data
        """
        payload = request if isinstance(request, dict) else {
            "messages": request.messages,
            "sessionId": request.session_id
        }
        response = self._request("POST", "/gateway/rehydrate", payload)
        return RehydrateResponse(
            original=response["original"],
            tokens_restored=response.get("tokensRestored", 0)
        )

    def health(self) -> Dict[str, str]:
        """
        Check if client is healthy.
        
        Returns:
            Health check result
            
        Example:
            >>> health = client.health()
            >>> print(health["status"])  # "ok"
        """
        return self._request("GET", "/health")

    def update_config(self, **kwargs: Any) -> None:
        """
        Update client configuration.
        
        Args:
            **kwargs: Configuration options to update
            
        Example:
            >>> client.update_config(tenant_id="new-tenant", timeout_ms=10000)
        """
        if "base_url" in kwargs:
            self._base_url = kwargs["base_url"].rstrip("/")
        if "token" in kwargs:
            self._token = kwargs["token"]
        if "api_key" in kwargs:
            self._api_key = kwargs["api_key"]
        if "tenant_id" in kwargs:
            self._tenant_id = kwargs["tenant_id"]
        if "user_id" in kwargs:
            self._user_id = kwargs["user_id"]
        if "timeout_ms" in kwargs:
            self._timeout_ms = kwargs["timeout_ms"]
        if "retries" in kwargs:
            self._retries = kwargs["retries"]
        if "retry_delay_ms" in kwargs:
            self._retry_delay_ms = kwargs["retry_delay_ms"]
        if "headers" in kwargs:
            self._headers = kwargs["headers"]

    def _request(
        self,
        method: str,
        path: str,
        payload: Optional[Any] = None
    ) -> Dict[str, Any]:
        """
        Make an HTTP request with retry logic.
        
        Args:
            method: HTTP method (GET, POST)
            path: Request path
            payload: Request payload (for POST)
            
        Returns:
            Response JSON
            
        Raises:
            GuardrailxError: On request failure
        """
        url = f"{self._base_url}{path}"
        headers = self._build_headers()
        timeout = self._timeout_ms / 1000

        attempt = 0
        last_error: Optional[Exception] = None

        while attempt <= self._retries:
            try:
                if method == "GET":
                    response = requests.get(url, headers=headers, timeout=timeout)
                else:  # POST
                    response = requests.post(
                        url,
                        headers=headers,
                        data=json.dumps(payload or {}),
                        timeout=timeout
                    )

                # Check for server errors (5xx) to retry
                if response.status_code >= 500:
                    raise GuardrailxError(
                        f"Server error {response.status_code}: {response.text}",
                        status=response.status_code,
                        response=response.text
                    )

                # Raise for other HTTP errors
                response.raise_for_status()

                return response.json()

            except requests.Timeout as exc:
                last_error = exc
                if attempt >= self._retries:
                    break
                time.sleep((self._retry_delay_ms / 1000) * (attempt + 1))
                attempt += 1

            except requests.RequestException as exc:
                status = getattr(exc.response, "status_code", None) if hasattr(exc, "response") else None
                text = getattr(exc.response, "text", str(exc)) if hasattr(exc, "response") else str(exc)
                
                # Retry on network errors or 5xx
                retryable = status is None or status >= 500
                
                if not retryable or attempt >= self._retries:
                    raise GuardrailxError(
                        f"GuardrailX error: {text}",
                        status=status,
                        response=text
                    ) from exc

                last_error = exc
                time.sleep((self._retry_delay_ms / 1000) * (attempt + 1))
                attempt += 1

        raise GuardrailxError(
            "GuardrailX request failed after retries",
            status=None,
            response=None
        ) from last_error

    def _build_headers(self) -> Dict[str, str]:
        """
        Build HTTP headers for requests.
        
        Returns:
            Headers dictionary
        """
        headers = {
            "content-type": "application/json",
            **self._headers,
        }

        # Prioritize API key authentication (like Node.js SDK)
        if self._api_key:
            headers["authorization"] = f"Bearer {self._api_key}"
        elif self._token:
            headers["authorization"] = (
                self._token if self._token.startswith("Bearer ") else f"Bearer {self._token}"
            )

        if self._tenant_id:
            headers["x-guardrailx-tenant"] = self._tenant_id

        if self._user_id:
            headers["x-guardrailx-user"] = self._user_id

        return headers

