"""
Risicare client for SDK initialization and configuration.

The client manages:
- SDK configuration (API key, endpoint, options)
- Context propagation patching
- Exporter lifecycle
- Provider instrumentation
"""

from __future__ import annotations

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

from risicare_core import patch_all, unpatch_all

logger = logging.getLogger("risicare")

from risicare.exporters.base import SpanExporter
from risicare.exporters.batch import BatchSpanProcessor

if TYPE_CHECKING:
    from risicare_core import Span


# =============================================================================
# Configuration
# =============================================================================

@dataclass
class RisicareConfig:
    """
    Configuration for the Risicare SDK.

    Attributes:
        api_key: API key for authentication with Risicare backend.
        endpoint: API endpoint URL.
        project_id: Optional project identifier.
        environment: Environment name (production, staging, development).
        service_name: Name of the service (for resource identification).
        service_version: Version of the service.
        enabled: Whether tracing is enabled.
        trace_content: Whether to capture prompt/completion content.
        sample_rate: Sampling rate (0.0 to 1.0, default 1.0 = trace all).
        batch_size: Number of spans per export batch.
        batch_timeout_ms: Maximum time to wait before flushing batch.
        auto_patch: Whether to automatically patch executors and asyncio.
        debug: Whether to enable debug logging.
    """
    api_key: Optional[str] = None
    endpoint: str = "https://app.risicare.ai"
    project_id: Optional[str] = None
    environment: str = "development"
    service_name: Optional[str] = None
    service_version: Optional[str] = None
    enabled: bool = True
    trace_content: bool = True
    sample_rate: float = 1.0
    batch_size: int = 100
    batch_timeout_ms: int = 1000
    auto_patch: bool = True
    debug: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)
    # OpenTelemetry configuration
    otlp_endpoint: Optional[str] = None
    otlp_headers: Optional[Dict[str, str]] = None
    otel_bridge: bool = False


# =============================================================================
# Client Singleton
# =============================================================================

_client: Optional["RisicareClient"] = None
_client_lock = threading.Lock()


def get_client() -> Optional["RisicareClient"]:
    """
    Get the global Risicare client instance.

    Returns:
        The client instance, or None if not initialized.
    """
    return _client


def init(
    api_key: Optional[str] = None,
    endpoint: Optional[str] = None,
    project_id: Optional[str] = None,
    environment: Optional[str] = None,
    service_name: Optional[str] = None,
    service_version: Optional[str] = None,
    enabled: Optional[bool] = None,
    trace_content: Optional[bool] = None,
    sample_rate: Optional[float] = None,
    batch_size: Optional[int] = None,
    batch_timeout_ms: Optional[int] = None,
    auto_patch: Optional[bool] = None,
    debug: Optional[bool] = None,
    exporters: Optional[List[SpanExporter]] = None,
    metadata: Optional[Dict[str, Any]] = None,
    otlp_endpoint: Optional[str] = None,
    otlp_headers: Optional[Dict[str, str]] = None,
    otel_bridge: Optional[bool] = None,
) -> "RisicareClient":
    """
    Initialize the Risicare SDK.

    This should be called once at application startup. It:
    - Configures the SDK with provided options
    - Patches ThreadPoolExecutor and asyncio for context propagation
    - Sets up span exporters
    - Registers atexit handler for graceful shutdown

    Configuration can be provided via:
    1. Function arguments (highest priority)
    2. Environment variables (RISICARE_*)
    3. Default values

    Args:
        api_key: API key (or RISICARE_API_KEY env var).
        endpoint: API endpoint (or RISICARE_ENDPOINT env var).
        project_id: Project ID (or RISICARE_PROJECT_ID env var).
        environment: Environment name (or RISICARE_ENVIRONMENT env var).
        service_name: Service name (or RISICARE_SERVICE_NAME env var).
        service_version: Service version (or RISICARE_SERVICE_VERSION env var).
        enabled: Enable tracing (or RISICARE_TRACING env var).
        trace_content: Capture content (or RISICARE_TRACE_CONTENT env var).
        sample_rate: Sampling rate 0.0-1.0 (or RISICARE_SAMPLE_RATE env var).
        batch_size: Spans per batch.
        batch_timeout_ms: Batch flush timeout.
        auto_patch: Auto-patch executors/asyncio.
        debug: Enable debug mode.
        exporters: Custom exporters (default: HTTP exporter if api_key provided).
        metadata: Additional metadata to attach to all spans.

    Returns:
        The initialized RisicareClient instance.

    Example:
        import risicare

        # Simple initialization
        risicare.init(api_key="rsk-...")

        # Full configuration
        risicare.init(
            api_key="rsk-...",
            environment="production",
            service_name="my-agent",
            sample_rate=0.1,  # Sample 10% of traces
            trace_content=False,  # Don't capture sensitive content
        )
    """
    global _client

    with _client_lock:
        if _client is not None:
            # Already initialized - update config if needed
            if debug or os.getenv("RISICARE_DEBUG"):
                logger.debug("SDK already initialized, returning existing client")
            return _client

        # Resolve sample rate
        resolved_sample_rate = sample_rate
        if resolved_sample_rate is None:
            env_rate = os.getenv("RISICARE_SAMPLE_RATE")
            if env_rate:
                try:
                    resolved_sample_rate = float(env_rate)
                except ValueError:
                    resolved_sample_rate = 1.0
            else:
                resolved_sample_rate = 1.0

        # Build configuration from arguments, env vars, and defaults
        config = RisicareConfig(
            api_key=api_key or os.getenv("RISICARE_API_KEY"),
            endpoint=endpoint or os.getenv("RISICARE_ENDPOINT", "https://app.risicare.ai"),
            project_id=project_id or os.getenv("RISICARE_PROJECT_ID"),
            environment=environment or os.getenv("RISICARE_ENVIRONMENT", "development"),
            service_name=service_name or os.getenv("RISICARE_SERVICE_NAME"),
            service_version=service_version or os.getenv("RISICARE_SERVICE_VERSION"),
            enabled=_resolve_bool(enabled, "RISICARE_TRACING", True),
            trace_content=_resolve_bool(trace_content, "RISICARE_TRACE_CONTENT", True),
            sample_rate=max(0.0, min(1.0, resolved_sample_rate)),  # Clamp to 0-1
            batch_size=batch_size or 100,
            batch_timeout_ms=batch_timeout_ms or 1000,
            auto_patch=auto_patch if auto_patch is not None else True,
            debug=_resolve_bool(debug, "RISICARE_DEBUG", False),
            metadata=metadata or {},
            otlp_endpoint=otlp_endpoint or os.getenv("RISICARE_OTLP_ENDPOINT"),
            otlp_headers=otlp_headers or _parse_headers_env("RISICARE_OTLP_HEADERS"),
            otel_bridge=_resolve_bool(otel_bridge, "RISICARE_OTEL_BRIDGE", False),
        )

        # Create client
        _client = RisicareClient(config, exporters=exporters)
        _client.start()

        # Register atexit handler for graceful shutdown
        atexit.register(_atexit_shutdown)

        return _client


def shutdown(timeout_ms: int = 5000) -> None:
    """
    Shutdown the Risicare SDK gracefully.

    Flushes pending spans and stops all background processing.
    Safe to call multiple times.

    Args:
        timeout_ms: Maximum time to wait for flush.

    Example:
        import risicare

        risicare.init(api_key="rsk-...")
        # ... your application code ...
        risicare.shutdown()  # Flush and cleanup before exit
    """
    global _client
    with _client_lock:
        if _client is not None:
            _client.shutdown(timeout_ms=timeout_ms)
            _client = None


def disable() -> None:
    """
    Temporarily disable tracing.

    Use this for A/B overhead measurement or when you need to
    temporarily stop span collection without shutting down the SDK.

    Paired with enable() to re-enable tracing.

    Example:
        import risicare

        risicare.init(api_key="rsk-...")

        # Measure with tracing
        with_tracing_time = measure_operation()

        # Measure without tracing
        risicare.disable()
        without_tracing_time = measure_operation()
        risicare.enable()

        overhead = with_tracing_time - without_tracing_time
    """
    global _client
    with _client_lock:
        if _client is not None:
            _client._config.enabled = False


def enable() -> None:
    """
    Re-enable tracing after disable().

    Must call init() first before enable() will have effect.
    """
    global _client
    with _client_lock:
        if _client is not None:
            _client._config.enabled = True


def is_enabled() -> bool:
    """
    Check if tracing is currently enabled.

    Returns:
        True if SDK is initialized and tracing is enabled, False otherwise.
    """
    global _client
    with _client_lock:
        return _client is not None and _client._config.enabled


def _atexit_shutdown() -> None:
    """Atexit handler for graceful shutdown."""
    try:
        shutdown(timeout_ms=2000)  # Shorter timeout for atexit
    except Exception:
        pass  # Ignore errors during atexit


def _reset_for_testing() -> None:
    """
    Reset all global state for testing.

    This is intended for test fixtures only. Do not use in production.
    """
    global _client
    with _client_lock:
        if _client is not None:
            try:
                _client.shutdown(timeout_ms=100)
            except Exception:
                pass
            _client = None


def _resolve_bool(value: Optional[bool], env_var: str, default: bool) -> bool:
    """Resolve boolean from argument, env var, or default."""
    if value is not None:
        return value
    env_value = os.getenv(env_var)
    if env_value is not None:
        return env_value.lower() in ("true", "1", "yes", "on")
    return default


def _parse_headers_env(env_var: str) -> Optional[Dict[str, str]]:
    """Parse headers from environment variable (comma-separated key=value pairs)."""
    raw = os.getenv(env_var)
    if not raw:
        return None
    headers: Dict[str, str] = {}
    for pair in raw.split(","):
        pair = pair.strip()
        if "=" in pair:
            key, value = pair.split("=", 1)
            headers[key.strip()] = value.strip()
    return headers or None


# =============================================================================
# Client Class
# =============================================================================

class RisicareClient:
    """
    Main Risicare client managing SDK lifecycle.

    Typically accessed via the global `init()` function rather than
    instantiated directly.
    """

    def __init__(
        self,
        config: RisicareConfig,
        exporters: Optional[List[SpanExporter]] = None,
    ):
        """
        Initialize the client.

        Args:
            config: SDK configuration.
            exporters: Optional list of span exporters.
        """
        self._config = config
        self._exporters: List[SpanExporter] = exporters or []
        self._processor: Optional[BatchSpanProcessor] = None
        self._started = False
        self._lock = threading.Lock()

        # Set up default exporter if API key provided and no custom exporters
        if not self._exporters and config.api_key:
            from risicare.exporters.http import HttpExporter
            self._exporters.append(HttpExporter(
                endpoint=config.endpoint,
                api_key=config.api_key,
                project_id=config.project_id,
                environment=config.environment,
            ))

        # Add OTLP exporter if endpoint configured
        if config.otlp_endpoint:
            from risicare.exporters.otlp import OTLPExporter
            self._exporters.append(OTLPExporter(
                endpoint=config.otlp_endpoint,
                headers=config.otlp_headers or {},
                service_name=config.service_name,
                service_version=config.service_version,
                environment=config.environment,
            ))

        # Add console exporter in debug mode
        if config.debug:
            from risicare.exporters.console import ConsoleExporter
            # Check if console exporter already added
            has_console = any(
                e.__class__.__name__ == "ConsoleExporter"
                for e in self._exporters
            )
            if not has_console:
                self._exporters.append(ConsoleExporter())

    @property
    def config(self) -> RisicareConfig:
        """Get the SDK configuration."""
        return self._config

    @property
    def is_enabled(self) -> bool:
        """Check if tracing is enabled."""
        return self._config.enabled and self._started

    def start(self) -> None:
        """
        Start the client.

        This patches executors/asyncio and starts the span processor.
        Called automatically by init().
        """
        with self._lock:
            if self._started:
                return

            if not self._config.enabled:
                if self._config.debug:
                    logger.debug("Tracing disabled, skipping initialization")
                return

            # Patch executors and asyncio for context propagation
            if self._config.auto_patch:
                patch_all()
                if self._config.debug:
                    logger.debug("Patched ThreadPoolExecutor and asyncio.create_task")

            # Install import hooks for auto-instrumentation of LLM providers
            from risicare.instrumentation import install_import_hooks
            install_import_hooks()

            # Start batch processor
            if self._exporters:
                self._processor = BatchSpanProcessor(
                    exporters=self._exporters,
                    batch_size=self._config.batch_size,
                    batch_timeout_ms=self._config.batch_timeout_ms,
                    debug=self._config.debug,
                )
                self._processor.start()
                if self._config.debug:
                    logger.debug("Started batch processor with %d exporter(s)", len(self._exporters))

            self._started = True
            if self._config.debug:
                logger.debug("SDK initialized successfully")

    def shutdown(self, timeout_ms: int = 5000) -> None:
        """
        Shutdown the client gracefully.

        Flushes pending spans and stops the processor.

        Args:
            timeout_ms: Maximum time to wait for flush.
        """
        with self._lock:
            if not self._started:
                return

            if self._config.debug:
                logger.debug("Shutting down...")

            # Stop processor (flushes pending spans)
            if self._processor:
                self._processor.shutdown(timeout_ms=timeout_ms)
                self._processor = None

            # Unpatch executors
            if self._config.auto_patch:
                unpatch_all()

            self._started = False
            if self._config.debug:
                logger.debug("Shutdown complete")

    def flush(self, timeout_ms: int = 5000) -> bool:
        """
        Flush pending spans.

        Args:
            timeout_ms: Maximum time to wait for flush.

        Returns:
            True if flush completed within timeout.
        """
        if self._processor:
            return self._processor.flush(timeout_ms=timeout_ms)
        return True

    def export_span(self, span: "Span") -> None:
        """
        Export a span to all configured exporters.

        Called by the tracer when a span ends.

        Args:
            span: The span to export.
        """
        if not self.is_enabled:
            return

        if self._processor:
            self._processor.on_span_end(span)

    def add_exporter(self, exporter: SpanExporter) -> None:
        """
        Add an exporter to the client.

        Args:
            exporter: The exporter to add.
        """
        self._exporters.append(exporter)
        if self._processor:
            self._processor.add_exporter(exporter)
