"""
PrivacyPal Python SDK Client.

A 1-to-1 mirror of the Node.js ``@privacypal/sdk`` ``PrivacyPalClient``.
Every public method, endpoint, payload shape, and error-message format
matches the JS SDK so that consuming code behaves identically regardless
of language.
"""

from __future__ import annotations

import logging
from typing import Any, Callable, Dict, List, Optional, Sequence

import requests

from .exceptions import (
    AuthenticationError,
    NetworkError,
    PrivacyPalError,
    RequestError,
    TrialExpiredError,
)

logger = logging.getLogger(__name__)

_DEFAULT_TIMEOUT = 30  # seconds
_FILE_UPLOAD_TIMEOUT = 120


class PrivacyPalClient:
    """Primary SDK client for all PrivacyPal API interactions.

    Parameters
    ----------
    api_url : str
        Base URL for the PrivacyPal API (e.g. ``http://localhost:42026``).
    api_key : str
        JWT token.  May be empty for ``login`` / ``register`` only.
    timeout : int, optional
        Default request timeout in seconds (default 30).
    on_token_refresh : callable, optional
        Invoked with the new token string whenever the caller refreshes it.
    """

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    def __init__(
        self,
        api_url: str,
        api_key: str = "",
        timeout: int = _DEFAULT_TIMEOUT,
        on_token_refresh: Optional[Callable[[str], None]] = None,
    ):
        self._api_url = api_url.rstrip("/")
        self._api_key = api_key
        self._timeout = timeout
        self._on_token_refresh = on_token_refresh
        self._session = requests.Session()
        self._sync_headers()

    # ------------------------------------------------------------------
    # Configuration helpers
    # ------------------------------------------------------------------

    def _sync_headers(self) -> None:
        self._session.headers.update({
            "Content-Type": "application/json",
        })
        if self._api_key:
            self._session.headers["x-access-token"] = self._api_key
            self._session.headers["Authorization"] = f"Bearer {self._api_key}"
        else:
            self._session.headers.pop("x-access-token", None)
            self._session.headers.pop("Authorization", None)

    def update_api_key(self, new_api_key: str) -> None:
        """Update the JWT token used for requests."""
        self._api_key = new_api_key
        self._sync_headers()
        if self._on_token_refresh:
            self._on_token_refresh(new_api_key)

    def update_api_url(self, new_api_url: str) -> None:
        """Update the API base URL."""
        self._api_url = new_api_url.rstrip("/")

    def get_config(self) -> Dict[str, Any]:
        """Return a read-only snapshot of the current configuration."""
        return {
            "apiUrl": self._api_url,
            "apiKey": self._api_key,
            "timeout": self._timeout,
        }

    # ------------------------------------------------------------------
    # Error handling (mirrors Node.js SDK error message patterns)
    # ------------------------------------------------------------------

    def _handle_error(self, error: requests.RequestException) -> None:
        """Translate *error* into the appropriate SDK exception.

        The message format matches the Node.js SDK exactly::

            "401: <message>"
            "403: <message> [Trial expired]"
            "Network Error: Unable to connect to PrivacyPal API at <url>"
            "Request Error: <message>"
        """
        resp = getattr(error, "response", None)
        if resp is not None:
            status = resp.status_code
            try:
                body = resp.json()
            except Exception:
                body = {}

            message = body.get("message") or body.get("error") or f"API Error: {status}"
            trial_expired = body.get("trialExpired") or body.get("requiresSubscription")
            trial_flag = " [Trial expired]" if trial_expired else ""

            if status == 401:
                raise AuthenticationError(f"401: {message}", status_code=401)
            if status == 403:
                raise TrialExpiredError(f"403: {message}{trial_flag}", status_code=403)
            raise PrivacyPalError(f"{message}", status_code=status)

        if isinstance(error, requests.ConnectionError):
            raise NetworkError(
                f"Network Error: Unable to connect to PrivacyPal API at {self._api_url}"
            )

        raise RequestError(f"Request Error: {error}")

    def _url(self, path: str) -> str:
        return f"{self._api_url}{path}"

    # ------------------------------------------------------------------
    # Encoding
    # ------------------------------------------------------------------

    def encode(
        self,
        data: str,
        *,
        source_container: str = "sdk_data",
        source_element: str = "text_input",
        metadata: Optional[Dict[str, Any]] = None,
        score_threshold: float = 0.35,
        language: str = "en",
        continuation_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Encode sensitive data by detecting PII and replacing with Privacy Twins.

        Returns a dict matching the Node.js ``EncodeResponse``::

            {
                "encodedData": "...",
                "continuationId": "...",
                "transformations": [...],
                "statistics": {...}
            }
        """
        payload: Dict[str, Any] = {
            "data": data,
            "sourceContainer": source_container,
            "sourceElement": source_element,
            "metadata": metadata or {},
            "scoreThreshold": score_threshold,
            "language": language,
        }
        if continuation_id:
            payload["continuationId"] = continuation_id

        try:
            resp = self._session.post(
                self._url("/api/scanner/encode"),
                json=payload,
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def encode_batch(
        self,
        items: Sequence[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Encode multiple data items in a single request.

        *items* is a list of dicts, each following the same shape as
        ``encode()`` parameters (``data``, ``source_container``, etc.).
        """
        payload = [
            {
                "data": item["data"],
                "sourceContainer": item.get("source_container", "sdk_data"),
                "sourceElement": item.get("source_element", "text_input"),
                "metadata": item.get("metadata", {}),
                "scoreThreshold": item.get("score_threshold", 0.35),
                "language": item.get("language", "en"),
            }
            for item in items
        ]
        try:
            resp = self._session.post(
                self._url("/api/scanner/encode/batch"),
                json=payload,
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def encode_file(
        self,
        file_content: bytes,
        file_name: str,
        *,
        process_images: bool = True,
        platform: Optional[str] = None,
        continuation_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Encode a file by uploading it for server-side PII detection.

        Supports PDF, DOCX, CSV, images, and more.
        """
        files = {"file": (file_name, file_content)}
        form_data: Dict[str, str] = {
            "processImages": "true" if process_images else "false",
        }
        if platform:
            form_data["platform"] = platform
        if continuation_id:
            form_data["continuationId"] = continuation_id

        headers = {}
        if self._api_key:
            headers["x-access-token"] = self._api_key
            headers["Authorization"] = f"Bearer {self._api_key}"

        try:
            resp = requests.post(
                self._url("/api/scanner/encode/file"),
                files=files,
                data=form_data,
                headers=headers,
                timeout=_FILE_UPLOAD_TIMEOUT,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Decoding
    # ------------------------------------------------------------------

    def decode(
        self,
        continuation_id: str,
        data: str,
        sensitive_hashes: Optional[List[str]] = None,
        authorization: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """Decode Privacy Twins back to original sensitive values.

        Returns a dict matching the Node.js ``DecodeResponse``.
        """
        auth = authorization or {
            "token": self._api_key,
            "purpose": "PrivacyPal SDK decoding",
            "type": "jwt",
        }
        payload = {
            "continuationId": continuation_id,
            "data": data,
            "sensitiveHashes": sensitive_hashes or [],
            "authorization": {
                "token": auth.get("token", self._api_key),
                "purpose": auth.get("purpose", "PrivacyPal SDK decoding"),
                "type": auth.get("type", "jwt"),
            },
        }
        try:
            resp = self._session.post(
                self._url("/api/scanner/decode"),
                json=payload,
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def get_dataset_twins(self, continuation_id: str) -> Dict[str, Any]:
        """Retrieve all Privacy Twins for a specific dataset."""
        try:
            resp = self._session.get(
                self._url(f"/api/scanner/twins/{continuation_id}"),
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # AI Chat
    # ------------------------------------------------------------------

    def chat_with_ai(self, **params: Any) -> Dict[str, Any]:
        """Send a prompt to an LLM with automatic PII encoding/decoding."""
        try:
            resp = self._session.post(
                self._url("/api/ai/chat"),
                json=params,
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def get_ai_providers(self) -> Dict[str, Any]:
        """Get available AI providers and models."""
        try:
            resp = self._session.get(
                self._url("/api/ai/providers"),
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Audit
    # ------------------------------------------------------------------

    def update_audit_tokens(
        self,
        continuation_id: str,
        *,
        tokens_in: int = 0,
        tokens_out: int = 0,
        model: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Update audit log entries with LLM token usage."""
        try:
            resp = self._session.post(
                self._url("/api/company/audit-logs/update-tokens"),
                json={
                    "continuationId": continuation_id,
                    "tokensIn": tokens_in,
                    "tokensOut": tokens_out,
                    "model": model,
                },
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return {"success": True, "data": resp.json()}
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Authentication (no token required for login / register)
    # ------------------------------------------------------------------

    def login(self, email: str, password: str) -> Dict[str, Any]:
        """Authenticate a user.  Does *not* require an existing API key.

        Returns the API response which includes ``data.token`` and user info.
        After calling this, pass the token to :meth:`update_api_key`.
        """
        try:
            resp = requests.post(
                self._url("/api/user/login"),
                json={"email": email, "password": password},
                headers={"Content-Type": "application/json"},
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def register(
        self,
        first_name: str,
        last_name: str,
        email: str,
        password: str,
    ) -> Dict[str, Any]:
        """Register a new user.  Does *not* require an existing API key.

        Returns the API response which includes ``token`` and ``newUserId``.
        """
        try:
            resp = requests.post(
                self._url("/api/user/register"),
                json={
                    "firstName": first_name,
                    "lastName": last_name,
                    "email": email,
                    "password": password,
                },
                headers={"Content-Type": "application/json"},
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def refresh_user_token(self, token: str) -> Dict[str, Any]:
        """Refresh an expired JWT token."""
        try:
            resp = self._session.post(
                self._url("/api/user/refresh-token"),
                json={"token": token},
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Account / User
    # ------------------------------------------------------------------

    def get_account(self) -> Dict[str, Any]:
        """Get the authenticated user's account information."""
        try:
            resp = self._session.get(
                self._url("/api/user/account"),
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def get_user_stats(self, days: Optional[int] = None) -> Dict[str, Any]:
        """Get user statistics.  Pass ``None`` for all-time."""
        path = f"/api/user/stats?days={days}" if days is not None else "/api/user/stats"
        try:
            resp = self._session.get(self._url(path), timeout=self._timeout)
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    def get_usage(self, days: int = 30) -> Dict[str, Any]:
        """Get usage data."""
        try:
            resp = self._session.get(
                self._url(f"/api/user/usage?days={days}"),
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Company
    # ------------------------------------------------------------------

    def get_company(self) -> Dict[str, Any]:
        """Get company information."""
        try:
            resp = self._session.get(self._url("/api/company"), timeout=self._timeout)
            resp.raise_for_status()
            data = resp.json()
            return data.get("company", data)
        except requests.RequestException as exc:
            self._handle_error(exc)

    def generate_invite_link(self) -> Dict[str, Any]:
        """Generate a team invite link."""
        try:
            resp = self._session.post(
                self._url("/api/company/invite"),
                timeout=self._timeout,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            self._handle_error(exc)

    # ------------------------------------------------------------------
    # Health
    # ------------------------------------------------------------------

    def health_check(self) -> Dict[str, Any]:
        """Check API connectivity.  Never raises — returns a status dict."""
        try:
            resp = self._session.get(self._url("/health"), timeout=5)
            return {
                "success": resp.status_code == 200,
                "status": resp.status_code,
                "data": resp.json() if resp.content else None,
            }
        except Exception as exc:
            return {"success": False, "error": str(exc)}
