"""
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 os
import logging
import threading
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)


@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})")

    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: str,
        response_text: Optional[str] = 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.

        Args:
            text: The input text to check for injection
            response_text: Optional agent response (for context in dashboard)
            source: Optional source identifier (e.g., 'customer_support_agent', 'search_agent')
            metadata: Optional dict with extra context (e.g., user_id, session_id)

        Example:
            shield.detect("user input here", response_text="agent response", source="support_agent")
        """
        thread = threading.Thread(
            target=self._send_ingest,
            args=(text, response_text, source, metadata),
            daemon=True,  # Won't prevent process exit
        )
        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}")

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

    def check(
        self,
        text: str,
        response_text: Optional[str] = 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).

        Args:
            text: The input text to check
            response_text: Optional agent response
            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."
            else:
                response = agent.run(user_input)
        """
        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/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}",
            )

