"""
Sentrial Client - Main SDK interface

Captures agent sessions, tool calls, and metrics. This data powers:
- Signal detection (auto-detect patterns/anomalies)
- Root cause analysis (understand WHY agents fail)
- Code fixer (AI-suggested fixes with GitHub PRs)

Dual-write mode with Langfuse:
    from sentrial import SentrialClient
    from sentrial.langfuse_integration import LangfuseIntegration
    
    client = SentrialClient(
        langfuse=LangfuseIntegration(public_key="pk-lf-...", secret_key="sk-lf-...")
    )
    
    # Data flows to both Sentrial AND Langfuse
    with client.begin(user_id="123", event="chat", input="Hello") as interaction:
        interaction.track_tool_call(...)
        interaction.set_output("World")
"""

import os
import logging
import requests
import threading
import time
from typing import Any, Optional, Dict, Union, TYPE_CHECKING
from .types import EventType
from .redact import PiiConfig, PiiBuiltinPatterns, PiiCustomPattern, redact_payload
from .batcher import EventBatcher, BatcherConfig

if TYPE_CHECKING:
    from .langfuse_integration import LangfuseIntegration, LangfuseTraceContext

# Retry configuration
MAX_RETRIES = 3
INITIAL_BACKOFF_SECONDS = 0.5
MAX_BACKOFF_SECONDS = 8.0
BACKOFF_MULTIPLIER = 2.0
RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}

# Module logger
_logger = logging.getLogger("sentrial")


class SentrialClient:
    """
    Sentrial Client for agent observability.
    
    Captures the data needed for automated root cause analysis and fix suggestions.

    Usage:
        client = SentrialClient(
            api_key="sentrial_live_xxx",
            api_url="https://api.sentrial.com"
        )

        # Create a session for an agent run
        session_id = client.create_session(
            name="Support Request #123",
            agent_name="support-agent"
        )

        # Track events
        client.track_tool_call(
            session_id=session_id,
            tool_name="search_kb",
            tool_input={"query": "password reset"},
            tool_output={"articles": ["KB-001"]}
        )

        # Complete session with metrics
        client.complete_session(
            session_id=session_id,
            success=True,
            custom_metrics={"customer_satisfaction": 90}
        )
    
    Note:
        By default, the SDK fails silently (fail_silently=True) to ensure
        monitoring never crashes your application. Set fail_silently=False
        during development to see errors.
    """

    def __init__(
        self,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
        fail_silently: bool = True,
        langfuse: Optional["LangfuseIntegration"] = None,
        pii: Optional[Union[PiiConfig, bool]] = None,
        batching: Optional[Union[BatcherConfig, bool]] = None,
    ):
        """
        Initialize Sentrial client.

        Args:
            api_url: URL of the Sentrial API server (defaults to SENTRIAL_API_URL env var or production)
            api_key: API key for authentication (defaults to SENTRIAL_API_KEY env var)
            fail_silently: If True (default), SDK errors are logged but won't crash your app.
                          Set to False during development to see full errors.
            langfuse: Optional LangfuseIntegration for dual-write mode. When provided,
                     all tracking calls are mirrored to Langfuse for side-by-side comparison.
            pii: PII redaction configuration. Pass a PiiConfig object for full control,
                 or True to fetch the org's PII config from the server automatically.
            batching: Event batching configuration. Pass True to enable with defaults,
                     or a BatcherConfig for full control. When enabled, fire-and-forget
                     tracking calls are queued and flushed periodically.
        """
        self.api_url = (api_url or os.environ.get("SENTRIAL_API_URL", "https://api.sentrial.com")).rstrip("/")
        self.api_key = api_key or os.environ.get("SENTRIAL_API_KEY")
        self.fail_silently = fail_silently
        self.langfuse = langfuse
        self._pii_config_needs_hydration = False
        self._pii_hydration_lock = threading.Lock()

        if pii is True:
            # Minimal default until we fetch the real config from the server
            self.pii = PiiConfig(enabled=True)
            self._pii_config_needs_hydration = True
        elif isinstance(pii, PiiConfig):
            self.pii = pii
        else:
            self.pii = pii

        self.session = requests.Session()
        # Thread-local storage for session state to prevent race conditions
        # in multi-threaded environments (e.g., FastAPI, multi-threaded agents)
        self._local = threading.local()

        if self.api_key:
            self.session.headers.update({"Authorization": f"Bearer {self.api_key}"})

        # Initialize event batcher if enabled
        self._batcher: Optional[EventBatcher] = None
        if batching is True:
            self._batcher = EventBatcher(
                send_fn=self._safe_request,
                config=BatcherConfig(enabled=True),
            )
        elif isinstance(batching, BatcherConfig) and batching.enabled:
            self._batcher = EventBatcher(
                send_fn=self._safe_request,
                config=batching,
            )

        # Per-session cost/token accumulator — populated by track_tool_call/track_decision
        self._session_accumulators: dict[str, dict[str, float]] = {}
        self._accumulator_lock = threading.Lock()

    @property
    def current_state(self) -> dict[str, Any]:
        """
        Thread-local state storage for the current session.
        
        Each thread gets its own isolated state dict, preventing
        race conditions when multiple threads use the same client.
        """
        if not hasattr(self._local, 'state'):
            self._local.state = {}
        return self._local.state
    
    @current_state.setter
    def current_state(self, value: dict[str, Any]) -> None:
        """Set the current state (used internally for clearing)."""
        self._local.state = value

    def _accumulate(self, session_id: str, cost: float = 0.0, token_count: Optional[int] = None, tool_output: Optional[dict] = None) -> None:
        """Accumulate cost/token data for a session from individual events."""
        with self._accumulator_lock:
            acc = self._session_accumulators.get(session_id)
            if acc is None:
                acc = {"cost": 0.0, "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
                self._session_accumulators[session_id] = acc
            if cost > 0:
                acc["cost"] += cost
            if token_count is not None and token_count > 0:
                acc["total_tokens"] += token_count
            # LLM wrappers store prompt/completion breakdown in tool_output["tokens"]
            tokens = tool_output.get("tokens") if isinstance(tool_output, dict) else None
            if isinstance(tokens, dict):
                acc["prompt_tokens"] += tokens.get("prompt", 0)
                acc["completion_tokens"] += tokens.get("completion", 0)

    def _hydrate_pii_config(self) -> None:
        """
        Fetch the organization's PII config from the server.

        Called lazily on the first request when ``pii=True`` was passed to the
        constructor.  Uses a threading lock so concurrent threads don't trigger
        duplicate fetches.
        """
        if not self._pii_config_needs_hydration:
            return

        with self._pii_hydration_lock:
            # Double-check after acquiring the lock
            if not self._pii_config_needs_hydration:
                return

            try:
                headers = {}
                if self.api_key:
                    headers["Authorization"] = f"Bearer {self.api_key}"

                resp = requests.get(
                    f"{self.api_url}/api/sdk/pii-config",
                    headers=headers,
                    timeout=10,
                )
                if resp.ok:
                    data = resp.json()
                    config = data.get("config")
                    if config:
                        builtin = config.get("builtinPatterns") or {}
                        custom_raw = config.get("customPatterns") or []

                        self.pii = PiiConfig(
                            enabled=config.get("enabled", True),
                            mode=config.get("mode", "label"),
                            fields=config.get("fields"),
                            enhanced_detection=config.get("enhancedDetection", False),
                            builtin_patterns=PiiBuiltinPatterns(
                                emails=builtin.get("emails", True),
                                phones=builtin.get("phones", True),
                                ssns=builtin.get("ssns", True),
                                credit_cards=builtin.get("creditCards", True),
                                ip_addresses=builtin.get("ipAddresses", True),
                            ),
                            custom_patterns=[
                                PiiCustomPattern(pattern=cp["pattern"], label=cp["label"])
                                for cp in custom_raw
                            ],
                        )
            except Exception:
                # Fail silently -- keep the minimal default config (enabled=True)
                _logger.debug("Sentrial: Failed to fetch PII config from server")

            self._pii_config_needs_hydration = False

    def close(self):
        """Close the underlying HTTP session, batcher, and Langfuse integration."""
        if self._batcher:
            self._batcher.shutdown()
        if self.session:
            self.session.close()
        if self.langfuse:
            self.langfuse.shutdown()

    def __enter__(self) -> "SentrialClient":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
        """Context manager exit."""
        self.close()
        return False

    def _safe_request(
        self,
        method: str,
        url: str,
        **kwargs,
    ) -> Optional[requests.Response]:
        """
        Make an HTTP request with graceful error handling and retry logic.
        
        Implements exponential backoff for transient failures (network errors,
        rate limits, server errors). Retries up to MAX_RETRIES times.
        
        If fail_silently is True, errors are logged and None is returned.
        If fail_silently is False, exceptions are re-raised.
        """
        # Hydrate server-driven PII config before first request
        if self._pii_config_needs_hydration:
            self._hydrate_pii_config()

        # Apply PII redaction to JSON payloads before sending
        if self.pii and "json" in kwargs and kwargs["json"]:
            kwargs["json"] = redact_payload(kwargs["json"], self.pii)

        last_exception: Optional[Exception] = None
        backoff = INITIAL_BACKOFF_SECONDS

        for attempt in range(MAX_RETRIES + 1):
            try:
                response = getattr(self.session, method)(url, **kwargs)
                
                # Check if we should retry based on status code
                if response.status_code in RETRYABLE_STATUS_CODES:
                    if attempt < MAX_RETRIES:
                        _logger.debug(
                            f"Sentrial: Retryable status {response.status_code} "
                            f"({method.upper()} {url}), retrying in {backoff:.1f}s (attempt {attempt + 1}/{MAX_RETRIES})"
                        )
                        time.sleep(backoff)
                        backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_SECONDS)
                        continue
                
                response.raise_for_status()
                return response
                
            except requests.exceptions.ConnectionError as e:
                # Network connectivity issues - always retry
                last_exception = e
                if attempt < MAX_RETRIES:
                    _logger.debug(
                        f"Sentrial: Connection error ({method.upper()} {url}), "
                        f"retrying in {backoff:.1f}s (attempt {attempt + 1}/{MAX_RETRIES})"
                    )
                    time.sleep(backoff)
                    backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_SECONDS)
                    continue
                    
            except requests.exceptions.Timeout as e:
                # Timeout - retry
                last_exception = e
                if attempt < MAX_RETRIES:
                    _logger.debug(
                        f"Sentrial: Timeout ({method.upper()} {url}), "
                        f"retrying in {backoff:.1f}s (attempt {attempt + 1}/{MAX_RETRIES})"
                    )
                    time.sleep(backoff)
                    backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_SECONDS)
                    continue
                    
            except requests.exceptions.RequestException as e:
                # Other request errors - don't retry
                last_exception = e
                break
        
        # All retries exhausted or non-retryable error
        if self.fail_silently:
            _logger.warning(f"Sentrial: Request failed ({method.upper()} {url}): {last_exception}")
            return None
        if last_exception:
            raise last_exception
        return None

    def create_session(
        self,
        name: str,
        agent_name: str,
        user_id: str,
        parent_session_id: Optional[str] = None,
        convo_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> Optional[str]:
        """
        Create a new session.

        Args:
            name: Name of the session
            agent_name: Required identifier for the agent type (used for grouping)
            user_id: Required external user ID (for grouping sessions by end-user)
            parent_session_id: Optional parent session ID for multi-agent hierarchies
            convo_id: Optional conversation ID to group related sessions into a thread
            metadata: Optional metadata

        Returns:
            Session ID, or None if the request failed and fail_silently is True
        """
        # Clear state from previous sessions to prevent memory leaks
        self.current_state.clear()

        payload: dict[str, Any] = {
            "name": name,
            "agentName": agent_name,
            "userId": user_id,
            "metadata": metadata,
        }

        if parent_session_id is not None:
            payload["parentSessionId"] = parent_session_id

        if convo_id is not None:
            payload["convoId"] = convo_id

        response = self._safe_request("post", f"{self.api_url}/api/sdk/sessions", json=payload)
        if response is None:
            return None

        try:
            data = response.json()
            return data["id"]
        except (ValueError, KeyError) as e:
            if self.fail_silently:
                _logger.warning(f"Sentrial: Failed to parse session response: {e}")
                return None
            raise

    def track_tool_call(
        self,
        session_id: str,
        tool_name: str,
        tool_input: dict[str, Any],
        tool_output: dict[str, Any],
        reasoning: Optional[str] = None,
        estimated_cost: float = 0.0,
        tool_error: Optional[dict[str, Any]] = None,
        token_count: Optional[int] = None,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """
        Track a tool call event.

        Args:
            session_id: Session ID
            tool_name: Name of the tool
            tool_input: Tool input data
            tool_output: Tool output data
            reasoning: Optional reasoning for why this tool was called
            estimated_cost: Estimated cost in USD for this tool call
            tool_error: Error details if the tool failed (e.g., {"message": "...", "type": "..."})
            token_count: Number of tokens used by this tool call (for LLM-based tools)
            trace_id: OpenTelemetry trace ID for distributed tracing
            span_id: OpenTelemetry span ID for distributed tracing
            metadata: Additional metadata to attach to this event

        Returns:
            Event data
        """
        self._accumulate(session_id, cost=estimated_cost, token_count=token_count, tool_output=tool_output)
        state_before = self.current_state.copy()

        # Update current state (simplified)
        self.current_state[f"{tool_name}_result"] = tool_output

        payload: dict[str, Any] = {
            "sessionId": session_id,
            "eventType": EventType.TOOL_CALL.value,
            "toolName": tool_name,
            "toolInput": tool_input,
            "toolOutput": tool_output,
            "reasoning": reasoning,
            "stateBefore": state_before,
            "stateAfter": self.current_state.copy(),
            "estimatedCost": estimated_cost,
        }
        
        if tool_error is not None:
            payload["toolError"] = tool_error
        if token_count is not None:
            payload["tokenCount"] = token_count
        if trace_id is not None:
            payload["traceId"] = trace_id
        if span_id is not None:
            payload["spanId"] = span_id
        if metadata is not None:
            payload["metadata"] = metadata

        if self._batcher:
            self._batcher.enqueue("post", f"{self.api_url}/api/sdk/events", json=payload)
            return {}
        response = self._safe_request("post", f"{self.api_url}/api/sdk/events", json=payload)
        if response is None:
            return {}
        return response.json()

    def track_decision(
        self,
        session_id: str,
        reasoning: str,
        alternatives: Optional[list[str]] = None,
        confidence: Optional[float] = None,
        estimated_cost: float = 0.0,
        token_count: Optional[int] = None,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """
        Track an LLM decision event.

        Args:
            session_id: Session ID
            reasoning: Decision reasoning
            alternatives: Alternative options considered
            confidence: Confidence score (0.0 to 1.0)
            estimated_cost: Estimated cost in USD for this decision
            token_count: Number of tokens used
            trace_id: OpenTelemetry trace ID for distributed tracing
            span_id: OpenTelemetry span ID for distributed tracing
            metadata: Additional metadata to attach to this event

        Returns:
            Event data
        """
        self._accumulate(session_id, cost=estimated_cost, token_count=token_count)
        state_before = self.current_state.copy()

        payload: dict[str, Any] = {
            "sessionId": session_id,
            "eventType": EventType.LLM_DECISION.value,
            "reasoning": reasoning,
            "alternativesConsidered": alternatives,
            "confidence": confidence,
            "stateBefore": state_before,
            "stateAfter": self.current_state.copy(),
            "estimatedCost": estimated_cost,
        }
        
        if token_count is not None:
            payload["tokenCount"] = token_count
        if trace_id is not None:
            payload["traceId"] = trace_id
        if span_id is not None:
            payload["spanId"] = span_id
        if metadata is not None:
            payload["metadata"] = metadata

        if self._batcher:
            self._batcher.enqueue("post", f"{self.api_url}/api/sdk/events", json=payload)
            return {}
        response = self._safe_request("post", f"{self.api_url}/api/sdk/events", json=payload)
        if response is None:
            return {}
        return response.json()

    def track_error(
        self,
        session_id: str,
        error_message: str,
        error_type: Optional[str] = None,
        tool_name: Optional[str] = None,
        stack_trace: Optional[str] = None,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """
        Track an error event.

        Args:
            session_id: Session ID
            error_message: Error message
            error_type: Type of error (e.g., "ValueError", "APIError")
            tool_name: Name of the tool that caused the error (if applicable)
            stack_trace: Stack trace for debugging
            trace_id: OpenTelemetry trace ID for distributed tracing
            span_id: OpenTelemetry span ID for distributed tracing
            metadata: Additional metadata to attach to this event

        Returns:
            Event data
        """
        state_before = self.current_state.copy()

        error_data: dict[str, Any] = {
            "message": error_message,
        }
        if error_type:
            error_data["type"] = error_type
        if stack_trace:
            error_data["stack_trace"] = stack_trace

        payload: dict[str, Any] = {
            "sessionId": session_id,
            "eventType": EventType.ERROR.value,
            "toolError": error_data,
            "stateBefore": state_before,
            "stateAfter": self.current_state.copy(),
        }
        
        if tool_name is not None:
            payload["toolName"] = tool_name
        if trace_id is not None:
            payload["traceId"] = trace_id
        if span_id is not None:
            payload["spanId"] = span_id
        if metadata is not None:
            payload["metadata"] = metadata

        if self._batcher:
            self._batcher.enqueue("post", f"{self.api_url}/api/sdk/events", json=payload)
            return {}
        response = self._safe_request("post", f"{self.api_url}/api/sdk/events", json=payload)
        if response is None:
            return {}
        return response.json()

    def update_state(self, key: str, value: Any):
        """Update the current state."""
        self.current_state[key] = value

    def complete_session(
        self, 
        session_id: str, 
        success: bool = True,
        failure_reason: Optional[str] = None,
        estimated_cost: Optional[float] = None,
        custom_metrics: Optional[Dict[str, float]] = None,
        duration_ms: Optional[int] = None,
        prompt_tokens: Optional[int] = None,
        completion_tokens: Optional[int] = None,
        total_tokens: Optional[int] = None,
        user_input: Optional[str] = None,
        assistant_output: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Complete a session with performance metrics.

        This is the recommended way to close sessions for performance monitoring.

        Args:
            session_id: Session ID
            success: Whether the session successfully completed its goal (default: True)
            failure_reason: If success=False, why did it fail?
            estimated_cost: Total estimated cost in USD for this session
            custom_metrics: Custom metrics (e.g., {"response_quality": 4.5, "steps_taken": 3})
            duration_ms: Duration in milliseconds (auto-calculated if not provided)
            prompt_tokens: Number of prompt/input tokens used
            completion_tokens: Number of completion/output tokens used
            total_tokens: Total tokens used (prompt + completion)
            user_input: The user's original query/input for this session
            assistant_output: The final assistant response for this session

        Returns:
            Updated session data

        Example:
            >>> client.complete_session(
            ...     session_id=session_id,
            ...     success=True,
            ...     estimated_cost=0.023,
            ...     prompt_tokens=1500,
            ...     completion_tokens=500,
            ...     total_tokens=2000,
            ...     user_input="What's the weather in San Francisco?",
            ...     assistant_output="The weather in San Francisco is 65°F and sunny.",
            ...     custom_metrics={
            ...         "customer_satisfaction": 4.5,
            ...         "order_value": 129.99,
            ...         "items_processed": 7
            ...     }
            ... )
        """
        # Flush queued events before completing the session so they arrive first
        if self._batcher:
            self._batcher.flush()

        # Auto-fill cost/tokens from accumulated event data if not explicitly provided
        with self._accumulator_lock:
            acc = self._session_accumulators.pop(session_id, None)
        if acc:
            if estimated_cost is None and acc["cost"] > 0:
                estimated_cost = acc["cost"]
            if prompt_tokens is None and acc["prompt_tokens"] > 0:
                prompt_tokens = int(acc["prompt_tokens"])
            if completion_tokens is None and acc["completion_tokens"] > 0:
                completion_tokens = int(acc["completion_tokens"])
            if total_tokens is None and acc["total_tokens"] > 0:
                total_tokens = int(acc["total_tokens"])

        payload = {
            "status": "completed" if success else "failed",
            "success": success,
        }

        if failure_reason is not None:
            payload["failureReason"] = failure_reason
        if estimated_cost is not None:
            payload["estimatedCost"] = estimated_cost
        if custom_metrics is not None:
            payload["customMetrics"] = custom_metrics
        if duration_ms is not None:
            payload["durationMs"] = duration_ms
        if prompt_tokens is not None:
            payload["promptTokens"] = prompt_tokens
        if completion_tokens is not None:
            payload["completionTokens"] = completion_tokens
        if total_tokens is not None:
            payload["totalTokens"] = total_tokens
        if user_input is not None:
            payload["userInput"] = user_input
        if assistant_output is not None:
            payload["assistantOutput"] = assistant_output
        
        response = self._safe_request("patch", f"{self.api_url}/api/sdk/sessions/{session_id}", json=payload)
        if response is None:
            return {}
        return response.json()

    def close_session(
        self,
        session_id: str,
        duration_ms: Optional[int] = None,
        success: Optional[bool] = None,
        failure_reason: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Mark a session as completed (legacy method).

        Note: Use complete_session() for better performance monitoring with custom metrics.

        Args:
            session_id: Session ID
            duration_ms: Duration in milliseconds
            success: Whether the session successfully completed its goal
            failure_reason: If success=False, why did it fail?

        Returns:
            Updated session data
        """
        return self.complete_session(
            session_id=session_id,
            success=success if success is not None else True,
            failure_reason=failure_reason,
            duration_ms=duration_ms,
        )

    # Cost calculation helpers
    @staticmethod
    def calculate_openai_cost(
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """
        Calculate cost for OpenAI API calls.

        Args:
            model: Model name (e.g., "gpt-4", "gpt-3.5-turbo")
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens

        Returns:
            Estimated cost in USD

        Example:
            >>> cost = SentrialClient.calculate_openai_cost(
            ...     model="gpt-4",
            ...     input_tokens=1000,
            ...     output_tokens=500
            ... )
            >>> print(f"Cost: ${cost:.4f}")
        """
        # Pricing as of 2026 (per 1M tokens)
        pricing = {
            "gpt-5.2": {"input": 5.0, "output": 15.0},
            "gpt-5": {"input": 4.0, "output": 12.0},
            "gpt-4.1": {"input": 2.0, "output": 8.0},
            "gpt-4.1-mini": {"input": 0.4, "output": 1.6},
            "gpt-4.1-nano": {"input": 0.1, "output": 0.4},
            "gpt-4o": {"input": 2.5, "output": 10.0},
            "gpt-4o-mini": {"input": 0.15, "output": 0.6},
            "gpt-4-turbo": {"input": 10.0, "output": 30.0},
            "gpt-4": {"input": 30.0, "output": 60.0},
            "gpt-3.5-turbo": {"input": 0.5, "output": 1.5},
            "o4-mini": {"input": 1.1, "output": 4.4},
            "o3": {"input": 10.0, "output": 40.0},
            "o3-mini": {"input": 1.1, "output": 4.4},
            "o3-pro": {"input": 20.0, "output": 80.0},
            "o1": {"input": 15.0, "output": 60.0},
            "o1-preview": {"input": 15.0, "output": 60.0},
            "o1-mini": {"input": 3.0, "output": 12.0},
        }

        # Find matching model (handle versioned models like gpt-4-0613)
        model_key = None
        for key in sorted(pricing.keys(), key=len, reverse=True):
            if model.startswith(key):
                model_key = key
                break

        if not model_key:
            # Default to gpt-4 pricing if unknown
            model_key = "gpt-4"

        rates = pricing[model_key]
        input_cost = (input_tokens / 1_000_000) * rates["input"]
        output_cost = (output_tokens / 1_000_000) * rates["output"]

        return input_cost + output_cost

    @staticmethod
    def calculate_anthropic_cost(
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """
        Calculate cost for Anthropic API calls.

        Args:
            model: Model name (e.g., "claude-3-opus", "claude-3-sonnet")
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens

        Returns:
            Estimated cost in USD
        """
        pricing = {
            "claude-opus-4": {"input": 15.0, "output": 75.0},
            "claude-sonnet-4": {"input": 3.0, "output": 15.0},
            "claude-4.5-opus": {"input": 20.0, "output": 100.0},
            "claude-4.5-sonnet": {"input": 4.0, "output": 20.0},
            "claude-4-opus": {"input": 18.0, "output": 90.0},
            "claude-4-sonnet": {"input": 3.5, "output": 17.5},
            "claude-3-7-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-5-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-5-haiku": {"input": 0.8, "output": 4.0},
            "claude-3-opus": {"input": 15.0, "output": 75.0},
            "claude-3-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-haiku": {"input": 0.25, "output": 1.25},
        }

        model_key = None
        for key in sorted(pricing.keys(), key=len, reverse=True):
            if model.startswith(key):
                model_key = key
                break

        if not model_key:
            model_key = "claude-3-sonnet"

        rates = pricing[model_key]
        input_cost = (input_tokens / 1_000_000) * rates["input"]
        output_cost = (output_tokens / 1_000_000) * rates["output"]

        return input_cost + output_cost

    @staticmethod
    def calculate_google_cost(
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """
        Calculate cost for Google/Gemini API calls.

        Args:
            model: Model name (e.g., "gemini-3-pro-preview", "gemini-2.5-flash")
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens

        Returns:
            Estimated cost in USD
        """
        # Pricing as of Jan 2026 (per 1M tokens)
        pricing = {
            # Gemini 3 - Preview (Jan 2026)
            "gemini-3-pro": {"input": 2.0, "output": 12.0},
            "gemini-3-flash": {"input": 0.5, "output": 3.0},
            # Gemini 2.5
            "gemini-2.5-pro": {"input": 1.25, "output": 5.0},
            "gemini-2.5-flash": {"input": 0.15, "output": 0.6},
            # Gemini 2.0
            "gemini-2.0-flash": {"input": 0.1, "output": 0.4},
            "gemini-2.0-flash-lite": {"input": 0.075, "output": 0.3},
            # Gemini 1.5
            "gemini-1.5-pro": {"input": 1.25, "output": 5.0},
            "gemini-1.5-flash": {"input": 0.075, "output": 0.3},
            # Gemini 1.0
            "gemini-1.0-pro": {"input": 0.5, "output": 1.5},
        }

        model_key = None
        for key in sorted(pricing.keys(), key=len, reverse=True):
            if model.startswith(key):
                model_key = key
                break

        if not model_key:
            model_key = "gemini-2.0-flash"

        rates = pricing[model_key]
        input_cost = (input_tokens / 1_000_000) * rates["input"]
        output_cost = (output_tokens / 1_000_000) * rates["output"]

        return input_cost + output_cost

    def flush(self) -> None:
        """Flush any queued events immediately. No-op if batching is not enabled."""
        if self._batcher:
            self._batcher.flush()

    def shutdown(self) -> None:
        """Shut down the event batcher, flushing remaining events. No-op if batching is not enabled."""
        if self._batcher:
            self._batcher.shutdown()

    def begin(
        self,
        user_id: str,
        event: str,
        input: Optional[str] = None,
        event_id: Optional[str] = None,
        convo_id: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> "Interaction":
        """
        Begin tracking an interaction (simplified API).

        Args:
            user_id: Required external user ID for grouping sessions
            event: Event type/name (e.g., "chat_message", "search_query")
            input: Optional input data for the interaction
            event_id: Optional custom event ID (auto-generated UUID if not provided)
            convo_id: Optional conversation ID to group related interactions
            metadata: Optional additional metadata

        Returns:
            Interaction object with finish() method

        Example:
            interaction = client.begin(
                user_id='user_123',
                event='chat_message',
                input=message,
                convo_id='convo_456'
            )

            # ... do your agent work ...

            interaction.finish(output=response_text)
        
        With context manager (recommended for auto-tracking):
            with client.begin(user_id='user_123', event='chat', input=message) as interaction:
                # Any wrapped LLM calls (wrap_openai, etc.) are auto-tracked!
                response = openai_client.chat.completions.create(...)
                interaction.set_output(response.choices[0].message.content)
                # Automatically finishes on exit
        
        With Langfuse dual-write (data flows to both systems):
            from sentrial.langfuse_integration import LangfuseIntegration
            client = SentrialClient(langfuse=LangfuseIntegration(...))
            
            with client.begin(user_id='user_123', event='chat', input=message) as interaction:
                # Data is mirrored to Langfuse automatically!
                interaction.track_tool_call(...)
                interaction.set_output(response)
        """
        import uuid as uuid_module  # noqa: local import to avoid top-level cost

        actual_event_id = event_id or str(uuid_module.uuid4())

        # Build metadata with optional fields
        full_metadata = metadata.copy() if metadata else {}
        if input is not None:
            full_metadata["input"] = input
        if convo_id:
            full_metadata["convo_id"] = convo_id
        full_metadata["event_id"] = actual_event_id

        # Create session first (clears old state)
        session_id = self.create_session(
            name=f"{event}:{actual_event_id[:8]}",
            agent_name=event,
            user_id=user_id,
            convo_id=convo_id,
            metadata=full_metadata,
        )

        # Store input in current_state AFTER create_session clears old state
        if input is not None:
            self.current_state["input"] = input

        # Create Langfuse trace context if integration is enabled
        langfuse_context: Optional["LangfuseTraceContext"] = None
        if self.langfuse and session_id:
            from .langfuse_integration import LangfuseTraceContext
            langfuse_context = LangfuseTraceContext(
                integration=self.langfuse,
                session_id=session_id,
                name=event,
                user_id=user_id,
                input_data=input,
                metadata=full_metadata,
            )

        interaction = Interaction(
            client=self,
            session_id=session_id,
            event_id=actual_event_id,
            user_id=user_id,
            event=event,
            langfuse_context=langfuse_context,
        )

        # Auto-wire session context so wrapped LLM calls flow into this session
        if session_id:
            from .wrappers import _current_session_id, _current_client
            interaction._session_token = _current_session_id.set(session_id)
            interaction._client_token = _current_client.set(self)

        return interaction


class Interaction:
    """
    Represents an in-progress interaction that can be finished.

    Created by SentrialClient.begin() - provides a clean begin/finish API pattern.
    
    Supports context manager for automatic cleanup:
        with client.begin(user_id="123", event="chat") as interaction:
            interaction.track_tool_call(...)
            # Automatically finishes on exit
    
    Auto-tracks wrapped LLM clients:
        When used as a context manager, any LLM calls made with wrapped clients
        (wrap_openai, wrap_anthropic, wrap_google) are automatically associated
        with this session.
    
    Langfuse dual-write:
        When the SentrialClient has a langfuse integration enabled, all tracking
        calls are automatically mirrored to Langfuse for side-by-side comparison.
    
    Note:
        If session creation failed (session_id is None), all tracking methods
        become no-ops to ensure your application continues running.
    """

    def __init__(
        self,
        client: SentrialClient,
        session_id: Optional[str],
        event_id: str,
        user_id: str,
        event: str,
        langfuse_context: Optional["LangfuseTraceContext"] = None,
    ):
        self.client = client
        self.session_id = session_id
        self.event_id = event_id
        self.user_id = user_id
        self.event = event
        self._start_time = time.time()
        self._finished = False
        self._success = True
        self._failure_reason: Optional[str] = None
        self._output: Optional[str] = None

        # Langfuse dual-write context
        self._langfuse_context = langfuse_context

        # Track if we're in a degraded state (session creation failed)
        self._degraded = session_id is None

        # Context var tokens for restoring previous session context (set by begin())
        self._session_token = None
        self._client_token = None

    def __enter__(self) -> "Interaction":
        """Context manager entry. Session context is already set by begin()."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
        """Context manager exit - auto-finishes the interaction and restores context."""
        if not self._finished:
            if exc_type is not None:
                self.finish(
                    success=False,
                    failure_reason=f"{exc_type.__name__}: {exc_val}" if exc_val else str(exc_type.__name__),
                )
            else:
                self.finish(
                    output=self._output,
                    success=self._success,
                    failure_reason=self._failure_reason,
                )
        return False  # Don't suppress exceptions
    
    def set_output(self, output: str) -> None:
        """
        Set the output for this interaction.
        
        This will be used when the context manager exits.
        Useful when you don't want to call finish() explicitly.
        
        Args:
            output: The output/response text
        """
        self._output = output

    def finish(
        self,
        output: Optional[str] = None,
        success: bool = True,
        failure_reason: Optional[str] = None,
        estimated_cost: Optional[float] = None,
        custom_metrics: Optional[Dict[str, float]] = None,
        duration_ms: Optional[int] = None,
        prompt_tokens: Optional[int] = None,
        completion_tokens: Optional[int] = None,
        total_tokens: Optional[int] = None,
    ) -> dict[str, Any]:
        """
        Finish the interaction and record metrics.

        Args:
            output: Output/response from the interaction
            success: Whether the interaction succeeded (default: True)
            failure_reason: Reason for failure if success=False
            estimated_cost: Total estimated cost in USD
            custom_metrics: Custom metrics
            prompt_tokens: Number of prompt/input tokens used
            completion_tokens: Number of completion/output tokens used
            total_tokens: Total tokens used

        Returns:
            Updated session data, or empty dict if degraded/already finished

        Example:
            interaction.finish(
                output="Here's the answer to your question...",
                success=True,
                custom_metrics={"satisfaction": 4.5}
            )
        """
        if self._finished:
            return {}  # Silently return instead of raising

        if self._degraded:
            self._finished = True
            return {}  # No-op if session creation failed

        self._finished = True

        # Restore previous session context (supports nesting)
        self._restore_session_context()

        # Use stored output if not provided
        final_output = output if output is not None else self._output

        # Merge output into custom_metrics if provided
        metrics = custom_metrics.copy() if custom_metrics else {}

        # Get input from session metadata if available
        user_input = self.client.current_state.get("input")

        # Finish Langfuse trace if enabled
        if self._langfuse_context:
            self._langfuse_context.finish(
                output=final_output,
                success=success,
                custom_metrics=metrics if metrics else None,
            )

        return self.client.complete_session(
            session_id=self.session_id,
            success=success,
            failure_reason=failure_reason,
            estimated_cost=estimated_cost,
            custom_metrics=metrics if metrics else None,
            duration_ms=duration_ms if duration_ms is not None else int((time.time() - self._start_time) * 1000),
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=total_tokens,
            user_input=user_input,
            assistant_output=final_output,
        )

    def _restore_session_context(self):
        """Restore the previous session context (set by begin())."""
        from .wrappers import _current_session_id, _current_client
        if self._session_token is not None:
            try:
                _current_session_id.reset(self._session_token)
            except ValueError:
                pass  # Token already consumed
            self._session_token = None
        if self._client_token is not None:
            try:
                _current_client.reset(self._client_token)
            except ValueError:
                pass  # Token already consumed
            self._client_token = None

    def track_tool_call(
        self,
        tool_name: str,
        tool_input: dict[str, Any],
        tool_output: dict[str, Any],
        reasoning: Optional[str] = None,
        estimated_cost: float = 0.0,
        tool_error: Optional[dict[str, Any]] = None,
        token_count: Optional[int] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """Track a tool call within this interaction."""
        if self._degraded:
            return {}  # No-op if session creation failed
        
        # Mirror to Langfuse if enabled
        if self._langfuse_context:
            self._langfuse_context.track_tool_call(
                tool_name=tool_name,
                tool_input=tool_input,
                tool_output=tool_output,
                tool_error=tool_error,
                metadata=metadata,
            )
        
        return self.client.track_tool_call(
            session_id=self.session_id,
            tool_name=tool_name,
            tool_input=tool_input,
            tool_output=tool_output,
            reasoning=reasoning,
            estimated_cost=estimated_cost,
            tool_error=tool_error,
            token_count=token_count,
            metadata=metadata,
        )

    def track_decision(
        self,
        reasoning: str,
        alternatives: Optional[list[str]] = None,
        confidence: Optional[float] = None,
        estimated_cost: float = 0.0,
        token_count: Optional[int] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """Track an LLM decision within this interaction."""
        if self._degraded:
            return {}  # No-op if session creation failed
        
        # Mirror to Langfuse as a generation if enabled
        if self._langfuse_context:
            self._langfuse_context.track_llm_call(
                name="llm_decision",
                input_data={"reasoning": reasoning, "alternatives": alternatives},
                output_data={"confidence": confidence},
                total_tokens=token_count,
                metadata=metadata,
            )
        
        return self.client.track_decision(
            session_id=self.session_id,
            reasoning=reasoning,
            alternatives=alternatives,
            confidence=confidence,
            estimated_cost=estimated_cost,
            token_count=token_count,
            metadata=metadata,
        )

    def track_error(
        self,
        error_message: str,
        error_type: Optional[str] = None,
        tool_name: Optional[str] = None,
        stack_trace: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """Track an error within this interaction."""
        # Store failure info for context manager exit
        self._success = False
        self._failure_reason = error_message
        
        if self._degraded:
            return {}  # No-op if session creation failed
        
        return self.client.track_error(
            session_id=self.session_id,
            error_message=error_message,
            error_type=error_type,
            tool_name=tool_name,
            stack_trace=stack_trace,
            metadata=metadata,
        )


# Module-level client instance for simple API
_default_client: Optional[SentrialClient] = None


def _get_client() -> SentrialClient:
    """Get or create the default client instance."""
    global _default_client
    if _default_client is None:
        _default_client = SentrialClient()
    return _default_client


def configure(
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    fail_silently: bool = True,
    pii: Optional[Union[PiiConfig, bool]] = None,
    batching: Optional[Union[BatcherConfig, bool]] = None,
) -> None:
    """
    Configure the default Sentrial client.

    Args:
        api_key: API key for authentication
        api_url: URL of the Sentrial API server
        fail_silently: If True (default), SDK errors won't crash your app
        pii: PII redaction configuration. Pass True to auto-fetch org config,
             or a PiiConfig object for full control.
        batching: Event batching configuration. Pass True to enable with defaults,
                 or a BatcherConfig for full control.
    """
    global _default_client
    _default_client = SentrialClient(
        api_key=api_key,
        api_url=api_url,
        fail_silently=fail_silently,
        pii=pii,
        batching=batching,
    )


def begin(
    user_id: str,
    event: str,
    input: Optional[str] = None,
    event_id: Optional[str] = None,
    convo_id: Optional[str] = None,
    metadata: Optional[dict[str, Any]] = None,
) -> Interaction:
    """
    Begin tracking an interaction (module-level convenience function).

    This is a shorthand for SentrialClient().begin(...).
    Configure the client first with sentrial.configure() or set SENTRIAL_API_KEY env var.

    Args:
        user_id: Required external user ID for grouping sessions
        event: Event type/name (e.g., "chat_message", "search_query")
        input: Optional input data for the interaction
        event_id: Optional custom event ID (auto-generated UUID if not provided)
        convo_id: Optional conversation ID to group related interactions
        metadata: Optional additional metadata

    Returns:
        Interaction object with finish() method

    Example:
        import sentrial

        sentrial.configure(api_key="sentrial_live_xxx")

        interaction = sentrial.begin(
            user_id='user_123',
            event='chat_message',
            input=message
        )

        # ... do your agent work ...

        interaction.finish(output=response_text)
    """
    return _get_client().begin(
        user_id=user_id,
        event=event,
        input=input,
        event_id=event_id,
        convo_id=convo_id,
        metadata=metadata,
    )

