"""
Cascade Security Shield SDK Client

Thin HTTP client for prompt injection detection.
Fire-and-forget by default — adds zero latency to your application.

Usage:
    from cascade.security import CascadeShield

    # Initialize once
    shield = CascadeShield(api_key="csk_live_...")

    # After your agent processes input and responds:
    user_input = "some user text"
    response = agent.run(user_input)
    send_to_user(response)

    # Fire-and-forget — non-blocking, runs in background thread
    shield.detect(user_input, response_text=response)

    # Or synchronous (for prevention mode — blocks until result):
    result = shield.check(user_input)
    if result.is_injection:
        print(f"Blocked: {result.reason}")
"""

import atexit
import json
import os
import logging
import threading
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

logger = logging.getLogger(__name__)

# Module-level registry of pending detect() threads. Used by atexit to flush
# before process exit so short scripts don't need time.sleep().
_pending_threads: List[threading.Thread] = []
_pending_lock = threading.Lock()


def _flush_pending_threads(timeout: float = 10.0) -> None:
    """Wait for all pending detect() threads to complete (called on process exit)."""
    with _pending_lock:
        threads = list(_pending_threads)
        _pending_threads.clear()
    for t in threads:
        t.join(timeout=timeout)


# Register atexit so short scripts (e.g. examples) don't need time.sleep(5)
atexit.register(_flush_pending_threads)


@dataclass
class ShieldResult:
    """Result from a synchronous security check."""
    is_injection: bool = False
    score: float = 0.0
    threshold: float = 0.5
    action: str = "allowed"
    reason: str = ""
    details: Dict[str, Any] = field(default_factory=dict)
    event_id: str = ""
    error: Optional[str] = None


class CascadeShield:
    """
    Cascade Security Shield client.
    
    Sends text to the Cascade backend for prompt injection detection.
    
    Two modes:
    - detect(): Fire-and-forget (async, non-blocking, zero latency)
    - check(): Synchronous (blocking, returns full result)
    
    Args:
        api_key: Cascade API key (or set CASCADE_API_KEY env var)
        endpoint: Backend URL (or set CASCADE_SHIELD_ENDPOINT env var).
                  Defaults to the CASCADE_ENDPOINT base URL.
        timeout: HTTP timeout in seconds (default 5)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        endpoint: Optional[str] = None,
        timeout: int = 5,
    ):
        self.api_key = api_key or os.getenv("CASCADE_API_KEY")
        if not self.api_key:
            logger.warning(
                "CascadeShield: No API key provided. "
                "Set api_key parameter or CASCADE_API_KEY environment variable."
            )

        # Derive endpoint: use explicit param, or SHIELD env, or strip /v1/traces from CASCADE_ENDPOINT
        if endpoint:
            self.endpoint = endpoint.rstrip("/")
        elif os.getenv("CASCADE_SHIELD_ENDPOINT"):
            self.endpoint = os.getenv("CASCADE_SHIELD_ENDPOINT").rstrip("/")
        else:
            # Derive from CASCADE_ENDPOINT (e.g., "https://api.runcascade.com/v1/traces" → "https://api.runcascade.com")
            cascade_ep = os.getenv("CASCADE_ENDPOINT", "https://api.runcascade.com/v1/traces")
            # Strip /v1/traces suffix to get the base URL
            base = cascade_ep.replace("/v1/traces", "").rstrip("/")
            self.endpoint = base

        self.timeout = timeout
        self._session = None

        logger.info(f"CascadeShield initialized (endpoint: {self.endpoint})")

    @staticmethod
    def _to_text(value: Any) -> str:
        """Convert any input to a text string for detection."""
        if isinstance(value, str):
            return value
        try:
            return json.dumps(value, default=str, ensure_ascii=False)
        except (TypeError, ValueError):
            return str(value)

    def _get_session(self):
        """Lazy-init requests session for connection pooling."""
        if self._session is None:
            import requests
            self._session = requests.Session()
            self._session.headers.update({
                "Content-Type": "application/json",
                "X-API-Key": self.api_key or "",
            })
        return self._session

    # ========================================================================
    # Fire-and-Forget Detection (async, non-blocking)
    # ========================================================================

    def detect(
        self,
        text: Union[str, Dict, List, Any],
        response_text: Union[str, Dict, List, Any, None] = None,
        source: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Send text for background injection detection. Non-blocking.

        This is the primary method for detection mode. It:
        1. Returns immediately (zero latency added to your app)
        2. Sends the text to Cascade in a background thread
        3. Cascade runs detection and logs the result
        4. Results appear in the Security Logs dashboard

        If the Cascade backend is unreachable, this fails silently.
        Your application is never impacted.

        Accepts strings, dicts, lists, or any object — non-string values
        are automatically serialized to JSON text for scanning.

        Args:
            text: The content to check for injection (str, dict, list, or any serializable object)
            response_text: Optional response content (str, dict, list, or any serializable object)
            source: Optional source identifier (e.g., 'hyperspell-retrieval', 'search_agent')
            metadata: Optional dict with extra context (e.g., user_id, session_id)

        Example:
            shield.detect("plain text input")
            shield.detect(response.documents, source="hyperspell-retrieval")
            shield.detect({"answer": "some text", "docs": [...]})
        """
        resolved_text = self._to_text(text)
        resolved_response = self._to_text(response_text) if response_text is not None else None
        thread = threading.Thread(
            target=self._send_ingest,
            args=(resolved_text, resolved_response, source, metadata),
            daemon=True,  # Won't prevent process exit (atexit flushes pending)
        )
        with _pending_lock:
            _pending_threads.append(thread)
        thread.start()

    def _send_ingest(
        self,
        text: str,
        response_text: Optional[str],
        source: Optional[str],
        metadata: Optional[Dict[str, Any]],
    ) -> None:
        """Background thread: send to /api/security/ingest."""
        try:
            session = self._get_session()
            payload = {"text": text}
            if response_text is not None:
                payload["response_text"] = response_text
            if source is not None:
                payload["source"] = source
            if metadata is not None:
                payload["metadata"] = metadata

            resp = session.post(
                f"{self.endpoint}/api/security/ingest",
                json=payload,
                timeout=self.timeout,
            )

            if resp.status_code == 200:
                data = resp.json()
                logger.debug(f"Shield detect: received=True, event_id={data.get('event_id')}")
            else:
                logger.warning(f"Shield detect failed: HTTP {resp.status_code} — {resp.text[:200]}")

        except Exception as e:
            # Silent failure — never impact the customer's app
            logger.debug(f"Shield detect error (silent): {e}")
        finally:
            # Remove from pending so long-running processes don't grow unbounded
            with _pending_lock:
                try:
                    _pending_threads.remove(threading.current_thread())
                except ValueError:
                    pass

    @staticmethod
    def flush(timeout: float = 10.0) -> None:
        """
        Wait for all pending detect() calls to complete.

        Call this before process exit if you need to ensure all detection
        requests have been sent. For short scripts, an atexit handler
        does this automatically — you typically don't need to call flush().

        Args:
            timeout: Max seconds to wait per pending thread (default 10)
        """
        _flush_pending_threads(timeout=timeout)

    # ========================================================================
    # Synchronous Check (for prevention mode)
    # ========================================================================

    def check(
        self,
        text: Union[str, Dict, List, Any],
        response_text: Union[str, Dict, List, Any, None] = None,
        source: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> ShieldResult:
        """
        Synchronous injection check. Blocks until result is returned.

        Use this for prevention mode — check BEFORE sending the response.
        Adds ~300-800ms latency (due to model inference).

        Accepts strings, dicts, lists, or any object — non-string values
        are automatically serialized to JSON text for scanning.

        Args:
            text: The content to check (str, dict, list, or any serializable object)
            response_text: Optional response content (str, dict, list, or any serializable object)
            source: Optional source identifier (e.g., 'customer_support_agent')
            metadata: Optional extra context

        Returns:
            ShieldResult with is_injection, score, reason, details

        Example:
            result = shield.check(user_input)
            if result.is_injection:
                return "Sorry, that input was flagged."
            result = shield.check(response.documents)
        """
        try:
            session = self._get_session()
            resolved_text = self._to_text(text)
            resolved_response = self._to_text(response_text) if response_text is not None else None
            payload = {"text": resolved_text}
            if resolved_response is not None:
                payload["response_text"] = resolved_response
            if source is not None:
                payload["source"] = source
            if metadata is not None:
                payload["metadata"] = metadata

            resp = session.post(
                f"{self.endpoint}/api/security/check",
                json=payload,
                timeout=self.timeout,
            )

            if resp.status_code == 200:
                data = resp.json()
                return ShieldResult(
                    is_injection=data.get("is_injection", False),
                    score=data.get("score", 0.0),
                    threshold=data.get("threshold", 0.5),
                    action=data.get("action", "allowed"),
                    reason=data.get("reason", ""),
                    details=data.get("details", {}),
                    event_id=data.get("event_id", ""),
                )
            else:
                return ShieldResult(
                    error=f"HTTP {resp.status_code}: {resp.text[:200]}",
                    reason="Check failed — backend error",
                )

        except Exception as e:
            logger.warning(f"Shield check error: {e}")
            return ShieldResult(
                error=str(e),
                reason=f"Check failed: {e}",
            )

