# Copyright (c) Metis. All rights reserved.

"""Initialization and lifecycle management for MantisDK tracing.

This module provides the core init(), shutdown(), and flush() functions
for managing the OpenTelemetry TracerProvider lifecycle.
"""

from __future__ import annotations

import logging
import threading
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

from opentelemetry import trace
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter

from .attributes import INSIGHT_ENVIRONMENT, INSIGHT_VERSION
from .exporters.insight import insight, is_insight_configured
from .processors import LoggingSpanProcessor, TracingContextSpanProcessor

if TYPE_CHECKING:
    from openinference.instrumentation import TraceConfig

logger = logging.getLogger(__name__)

# Module-level state
_initialized = False
_init_lock = threading.Lock()
_provider: Optional[TracerProvider] = None
_processors: List[BatchSpanProcessor] = []
_export_count = 0
_batch_config: Dict[str, Any] = {}  # Batch processor kwargs, set by init()


def init(
    *,
    trace_name: Optional[str] = None,
    service_name: Optional[str] = None,
    environment: Optional[str] = None,
    version: Optional[str] = None,
    # Insight configuration (alternative to env vars)
    insight_host: Optional[str] = None,
    insight_public_key: Optional[str] = None,
    insight_secret_key: Optional[str] = None,
    insight_endpoint: Optional[str] = None,
    instrument_default: bool = True,
    exporters: Optional[List[SpanExporter]] = None,
    trace_config: Optional["TraceConfig"] = None,
    force: bool = False,
    integrations: Optional[Dict[str, Any]] = None,
    debug: Union[bool, str] = False,
    # Batch processor tuning (passed through to BatchSpanProcessor)
    batch_queue_size: Optional[int] = None,
    batch_flush_interval: Optional[int] = None,
    batch_max_export_size: Optional[int] = None,
) -> None:
    """Initialize MantisDK standalone tracing.

    This function configures OpenTelemetry with a connection to Mantis Insight
    (if configured via environment variables), and automatically instruments
    installed AI libraries (OpenAI, Anthropic, LiteLLM, OpenAI Agents SDK, etc.).

    Provider Reuse Behavior:
        - If no TracerProvider exists, creates one.
        - If a TracerProvider exists and force=False, reuses it and adds
          processors/exporters (avoiding duplicates).
        - If force=True, replaces the existing TracerProvider.

    Args:
        trace_name: Optional name for the trace.
        service_name: Name of the service (mapped to resource attributes).
            Defaults to "mantisdk-tracing".
        environment: Environment name (e.g., "production", "staging", "development").
            Sets `insight.environment` resource attribute for filtering traces in Insight.
        version: Application version (e.g., "1.2.3", git SHA). Sets `insight.version`
            resource attribute for tracking regressions across versions.
        insight_host: Insight server URL (e.g., "https://insight.withmetis.ai").
            Overrides INSIGHT_HOST environment variable.
        insight_public_key: Insight public key (pk-lf-...).
            Overrides INSIGHT_PUBLIC_KEY environment variable.
        insight_secret_key: Insight secret key (sk-lf-...).
            Overrides INSIGHT_SECRET_KEY environment variable.
        insight_endpoint: Optional OTLP endpoint URL. Overrides INSIGHT_OTLP_ENDPOINT.
            If not provided, derived from insight_host.
        instrument_default: Whether to automatically instrument the "core set"
            of libraries (OpenAI Agents SDK, OpenAI, Anthropic, LiteLLM, etc.)
            if installed. Defaults to True.
        exporters: List of custom OpenTelemetry SpanExporters. If None and
            Insight credentials are configured (via params or env vars), auto-creates
            an Insight exporter.
        trace_config: OpenInference TraceConfig for controlling capture behavior
            (hiding inputs/outputs, etc.). Default captures all.
        force: If True, replaces any existing global TracerProvider. If False
            (default), reuses the existing provider and appends processors.
        integrations: Dictionary for fine-grained control over integrations.
        debug: Controls console logging of tracing activity. Accepts:
            - False (default): Silent operation, no console output.
            - True: Shortcut for "DEBUG" -- full firehose including every
              span start/end, export results, and internal plumbing.
            - "DEBUG": Same as True.
            - "INFO": Recommended. Shows init summary, export success/failure,
              instrumentor activation, and shutdown summary.
            - "WARNING": Minimal. Only shows export failures and missing
              exporter warnings.
            - "ERROR": Only export failures and exceptions.
        batch_queue_size: Max spans buffered in memory by BatchSpanProcessor.
            Default: 2048. Increase for bulk import workloads to prevent
            silent span drops. Wraps OTEL_BSP_MAX_QUEUE_SIZE.
        batch_flush_interval: How often (in ms) the processor flushes spans.
            Default: 5000. Decrease for lower-latency export.
            Wraps OTEL_BSP_SCHEDULE_DELAY.
        batch_max_export_size: Max spans per export call. Default: 512.
            Wraps OTEL_BSP_MAX_EXPORT_BATCH_SIZE.

    Example::

        import mantisdk.tracing as tracing

        # Minimal: auto-detect Insight from env vars, auto-instrument
        tracing.init()

        # With environment and version for filtering in Insight
        tracing.init(
            service_name="my-agent",
            environment="production",
            version="1.2.3",
            instrument_default=True,
        )

        # Explicit Insight credentials (instead of env vars)
        tracing.init(
            service_name="my-agent",
            environment="production",
            version="1.2.3",
            insight_host="https://insight.withmetis.ai",
            insight_public_key="pk-lf-...",
            insight_secret_key="sk-lf-...",
        )

        # Ensure cleanup on shutdown
        import atexit
        atexit.register(tracing.shutdown)
    """
    global _initialized, _provider

    resolved_level = _resolve_debug_level(debug)
    if resolved_level is not None:
        _configure_debug_logging(resolved_level)

    with _init_lock:
        # Check existing provider
        current_provider = trace.get_tracer_provider()
        has_sdk_provider = isinstance(current_provider, TracerProvider)

        if _initialized and not force:
            logger.debug("MantisDK tracing already initialized. Use force=True to reinitialize.")
            return

        if has_sdk_provider and not force:
            # Reuse existing provider
            logger.info("Reusing existing TracerProvider")
            _provider = current_provider
        else:
            # Create new provider
            resource_attrs = {
                SERVICE_NAME: service_name or "mantisdk-tracing",
            }
            if trace_name:
                resource_attrs["trace.name"] = trace_name
            if environment:
                resource_attrs[INSIGHT_ENVIRONMENT] = environment
            if version:
                resource_attrs[INSIGHT_VERSION] = version

            resource = Resource.create(resource_attrs)
            _provider = TracerProvider(resource=resource)

            # Set as global provider
            trace.set_tracer_provider(_provider)
            logger.info("Created new TracerProvider with service_name=%s", service_name or "mantisdk-tracing")

        # Add TracingContextSpanProcessor for automatic context injection
        # This must be added BEFORE exporter processors so context is set first
        _add_context_processor_idempotent(_provider)

        if resolved_level is not None:
            _add_logging_processor_idempotent(_provider)

        configured_exporters = exporters or []

        # Auto-configure Insight exporter if no custom exporters provided
        if not configured_exporters:
            # Check if Insight credentials are provided via parameters
            if insight_host and insight_public_key and insight_secret_key:
                try:
                    insight_exporter = insight(
                        host=insight_host,
                        public_key=insight_public_key,
                        secret_key=insight_secret_key,
                        endpoint=insight_endpoint,
                    )
                    configured_exporters.append(insight_exporter)
                    logger.info("Configured Insight exporter from init() parameters")
                except ValueError as e:
                    logger.warning("Could not configure Insight exporter: %s", e)
            # Fall back to environment variables
            elif is_insight_configured():
                try:
                    insight_exporter = insight()
                    configured_exporters.append(insight_exporter)
                    logger.info("Auto-configured Insight exporter from environment variables")
                except ValueError as e:
                    logger.warning("Could not auto-configure Insight exporter: %s", e)

        if not configured_exporters:
            logger.warning("No exporter configured -- spans will not be persisted")

        # Store batch processor config for _add_processor_idempotent
        _batch_config.clear()
        if batch_queue_size is not None:
            _batch_config["max_queue_size"] = batch_queue_size
        if batch_flush_interval is not None:
            _batch_config["schedule_delay_millis"] = batch_flush_interval
        if batch_max_export_size is not None:
            _batch_config["max_export_batch_size"] = batch_max_export_size

        # Add processors for each exporter (idempotent)
        for exporter in configured_exporters:
            _add_processor_idempotent(_provider, exporter)

        # Run instrumentation
        if instrument_default:
            instrument()

        _initialized = True
        logger.info("MantisDK tracing initialized successfully")


def instrument(
    names: Optional[List[str]] = None,
    skip: Optional[List[str]] = None,
) -> None:
    """Manually enable specific instrumentations.

    This function looks up instrumentors by name from the registry and
    activates them. If no names are provided and init() was called with
    instrument_default=True, this enables the "core set".

    Core Set (enabled by default):
        - openai
        - anthropic
        - langchain
        - llama_index
        - litellm

    Additional Available:
        - google_adk
        - mistral
        - groq
        - bedrock

    Args:
        names: List of instrumentor names to enable. If None and called from
            init() with instrument_default=True, enables the core set.
        skip: List of instrumentor names to explicitly skip.

    Example::

        import mantisdk.tracing_claude as tracing

        # Enable only specific instrumentors
        tracing.init(instrument_default=False)
        tracing.instrument(names=["openai", "anthropic"])

        # Enable all except langchain
        tracing.instrument(skip=["langchain"])
    """
    from .instrumentors.registry import get_registry

    registry = get_registry()
    skip_set = set(skip or [])

    # Determine which instrumentors to enable
    if names is None:
        # Default core set -- instruments all installed AI libraries automatically
        target_names = [
            "openai_agents",
            "claude_agent_sdk",
            "openai",
            "anthropic",
            "langchain",
            "llama_index",
            "litellm",
        ]
    else:
        target_names = names

    # Filter out skipped instrumentors
    target_names = [name for name in target_names if name not in skip_set]

    # Activate each instrumentor
    for name in target_names:
        instrumentor = registry.get(name)
        if instrumentor is not None:
            try:
                instrumentor.instrument()
                logger.info("Instrumented: %s", name)
            except Exception as e:
                logger.debug("Could not instrument %s: %s", name, e)
        else:
            logger.debug("Instrumentor not available (not installed?): %s", name)


def shutdown(timeout_millis: int = 30000) -> None:
    """Force flush all pending spans and shutdown the tracer provider.

    This function should be called before process exit to ensure all spans
    are exported. It is safe to call multiple times.

    Args:
        timeout_millis: Maximum time to wait for flush completion in milliseconds.

    Example::

        import mantisdk.tracing_claude as tracing
        import atexit

        tracing.init()
        atexit.register(tracing.shutdown)

        # ... run application ...
    """
    global _initialized, _provider, _context_processor_added, _logging_processor_added, _export_count

    with _init_lock:
        if not _initialized or _provider is None:
            logger.debug("Tracing not initialized, nothing to shutdown")
            return

        try:
            # Shutdown the provider (which flushes and shuts down processors)
            _provider.shutdown()
            logger.info("Shutdown complete: exported %d span(s) total", _export_count)
        except Exception as e:
            logger.warning("Error during TracerProvider shutdown: %s", e)

        _processors.clear()
        _batch_config.clear()
        _initialized = False
        _provider = None
        _context_processor_added = False
        _logging_processor_added = False
        _export_count = 0


def flush(timeout_millis: int = 30000) -> bool:
    """Force flush all pending spans without shutting down.

    This is useful for ensuring spans are exported at specific points
    (e.g., after completing a batch of work) without ending the tracing session.

    Args:
        timeout_millis: Maximum time to wait for flush completion in milliseconds.

    Returns:
        True if flush completed successfully, False otherwise.

    Example::

        import mantisdk.tracing_claude as tracing

        tracing.init()

        # Process batch
        for item in batch:
            process(item)

        # Ensure all spans from this batch are exported
        tracing.flush()
    """
    with _init_lock:
        if _provider is None:
            logger.debug("No provider to flush")
            return True

        try:
            success = _provider.force_flush(timeout_millis=timeout_millis)
            if success:
                logger.debug("Flush completed successfully")
            else:
                logger.warning("Flush timed out after %d ms", timeout_millis)
            return success
        except Exception as e:
            logger.warning("Error during flush: %s", e)
            return False


# Track if context processor has been added
_context_processor_added = False


def _add_context_processor_idempotent(provider: TracerProvider) -> None:
    """Add the TracingContextSpanProcessor if not already added.

    This processor injects run context (run_id, session_id) and tracing
    context (call_type, tags, environment) into all spans automatically.
    """
    global _context_processor_added

    if _context_processor_added:
        logger.debug("TracingContextSpanProcessor already added")
        return

    processor = TracingContextSpanProcessor()
    provider.add_span_processor(processor)
    _context_processor_added = True
    logger.debug("Added TracingContextSpanProcessor for automatic context injection")


def _add_processor_idempotent(provider: TracerProvider, exporter: SpanExporter) -> None:
    """Add a BatchSpanProcessor for the exporter, avoiding duplicates.

    Checks if a processor for the same exporter type/endpoint already exists
    to prevent duplicate exports.
    """
    # Check for existing processor with same exporter type
    exporter_id = _get_exporter_id(exporter)
    for existing in _processors:
        if _get_exporter_id(existing.span_exporter) == exporter_id:
            logger.debug("Processor already exists for exporter: %s", exporter_id)
            return

    # Create and add new processor.
    # Wrap Insight exporters with AnnotationFilterSpanProcessor to suppress
    # mantisdk.annotation reward spans from OTLP (rewards are sent as
    # Langfuse Scores via the API instead).
    from .exporters.insight import InsightOTLPExporter
    from .processors import AnnotationFilterSpanProcessor

    batch_processor = BatchSpanProcessor(exporter, **_batch_config)
    if isinstance(exporter, InsightOTLPExporter):
        filtered = AnnotationFilterSpanProcessor(batch_processor)
        provider.add_span_processor(filtered)
        logger.debug("Added filtered BatchSpanProcessor for: %s", exporter_id)
    else:
        provider.add_span_processor(batch_processor)
        logger.debug("Added BatchSpanProcessor for: %s", exporter_id)
    _processors.append(batch_processor)


def _get_exporter_id(exporter: SpanExporter) -> str:
    """Generate a unique identifier for an exporter instance.

    Uses class name + endpoint (if available) to identify exporters.
    """
    class_name = exporter.__class__.__name__

    # Try to get endpoint for OTLP exporters
    endpoint = getattr(exporter, "_endpoint", None)
    if endpoint:
        return f"{class_name}:{endpoint}"

    return class_name


def _resolve_debug_level(debug: Union[bool, str]) -> Optional[str]:
    """Resolve the debug parameter to a Python logging level name, or None if disabled."""
    if debug is False or debug is None:
        return None
    if debug is True:
        return "DEBUG"
    level_name = str(debug).upper()
    if level_name not in ("DEBUG", "INFO", "WARNING", "ERROR"):
        raise ValueError(f"Invalid debug level: {debug!r}. " "Use True, False, 'DEBUG', 'INFO', 'WARNING', or 'ERROR'.")
    return level_name


def _record_exported(count: int) -> None:
    """Increment the global export counter (called by InsightOTLPExporter)."""
    global _export_count
    _export_count += count


_logging_processor_added = False


def _add_logging_processor_idempotent(provider: TracerProvider) -> None:
    """Add the LoggingSpanProcessor if not already added."""
    global _logging_processor_added

    if _logging_processor_added:
        return

    provider.add_span_processor(LoggingSpanProcessor())
    _logging_processor_added = True
    logger.debug("Added LoggingSpanProcessor for debug visibility")


def _configure_debug_logging(level: str = "DEBUG") -> None:
    """Configure Python logging so mantisdk.tracing messages reach the console.

    Sets up a StreamHandler on the ``mantisdk.tracing`` logger hierarchy
    (and the OTLP exporter logger) at the requested level.
    Writes to stdout (not stderr) so output interleaves naturally with
    application print() statements.

    Args:
        level: Python logging level name (DEBUG, INFO, WARNING, ERROR).
    """
    import sys

    numeric_level = getattr(logging, level, logging.DEBUG)

    tracing_logger = logging.getLogger("mantisdk.tracing")
    if tracing_logger.handlers:
        return

    handler = logging.StreamHandler(stream=sys.stdout)
    handler.setLevel(numeric_level)
    fmt = logging.Formatter("[%(name)s] %(message)s")
    handler.setFormatter(fmt)

    tracing_logger.addHandler(handler)
    tracing_logger.setLevel(numeric_level)
    tracing_logger.propagate = False

    # Also capture OTLP exporter warnings (connection failures, retries)
    otel_logger = logging.getLogger("opentelemetry.exporter.otlp")
    otel_logger.addHandler(handler)
    otel_logger.setLevel(logging.WARNING)
    otel_logger.propagate = False

    # Suppress noisy OTel warnings about provider overriding (force=True triggers these)
    for name in ("opentelemetry.trace", "opentelemetry.sdk.trace"):
        lg = logging.getLogger(name)
        lg.setLevel(logging.ERROR)


def get_tracer(name: str = "mantisdk.tracing") -> trace.Tracer:
    """Get a tracer instance for creating spans.

    This is primarily for internal use. Users should prefer the
    trace() and span() context managers.

    Args:
        name: The tracer name (instrumentation scope).

    Returns:
        OpenTelemetry Tracer instance.
    """
    return trace.get_tracer(name)
