"""Initialization logic for the Flowlines SDK."""

from __future__ import annotations

import atexit
import importlib
import importlib.util
import logging
import threading
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

from opentelemetry import trace
from opentelemetry.sdk.trace import SpanProcessor, TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

from flowlines._context import FlowlinesSpanProcessor
from flowlines._context import clear_context as _clear_context
from flowlines._context import set_context as _set_context
from flowlines._exporter import FlowlinesExporter
from flowlines._patches import apply_all as _apply_patches

if TYPE_CHECKING:
    from collections.abc import Generator
    from contextvars import Token
    from types import ModuleType

logger = logging.getLogger("flowlines")

_DEFAULT_ENDPOINT = "https://ingest.flowlines.ai"

# Registry of instrumentors: (library_to_check, instrumentor_module, class_name)
_INSTRUMENTOR_REGISTRY: list[tuple[str, str, str]] = [
    ("openai", "opentelemetry.instrumentation.openai", "OpenAIInstrumentor"),
    ("anthropic", "opentelemetry.instrumentation.anthropic", "AnthropicInstrumentor"),
    (
        "google.genai",
        "opentelemetry.instrumentation.google_generativeai",
        "GoogleGenerativeAiInstrumentor",
    ),
    ("botocore", "opentelemetry.instrumentation.bedrock", "BedrockInstrumentor"),
    ("cohere", "opentelemetry.instrumentation.cohere", "CohereInstrumentor"),
    ("vertexai", "opentelemetry.instrumentation.vertexai", "VertexAIInstrumentor"),
    ("pinecone", "opentelemetry.instrumentation.pinecone", "PineconeInstrumentor"),
    ("together", "opentelemetry.instrumentation.together", "TogetherAiInstrumentor"),
    ("chromadb", "opentelemetry.instrumentation.chromadb", "ChromaInstrumentor"),
    ("qdrant_client", "opentelemetry.instrumentation.qdrant", "QdrantInstrumentor"),
    ("langchain", "opentelemetry.instrumentation.langchain", "LangchainInstrumentor"),
    (
        "llama_index",
        "opentelemetry.instrumentation.llamaindex",
        "LlamaIndexInstrumentor",
    ),
    ("mcp", "opentelemetry.instrumentation.mcp", "McpInstrumentor"),
]


_LOOPBACK_HOSTS = frozenset(("localhost", "127.0.0.1", "::1"))


def _validate_endpoint(endpoint: str) -> None:
    """Validate that the endpoint uses HTTPS unless it targets localhost.

    Raises:
        ValueError: If the endpoint is not HTTPS and not a loopback address.
    """
    parsed = urlparse(endpoint)
    hostname = parsed.hostname or ""
    if parsed.scheme == "https":
        return
    if hostname in _LOOPBACK_HOSTS:
        return
    msg = (
        f"Endpoint must use HTTPS unless targeting a loopback address. Got: {endpoint}"
    )
    raise ValueError(msg)


def _load_instrumentor(module_path: str, class_name: str) -> Any:
    """Dynamically load an instrumentor class and return an instance."""
    module: ModuleType = importlib.import_module(module_path)
    cls: type[Any] = getattr(module, class_name)
    return cls()


def _configure_verbose_logging() -> None:
    """Configure the ``flowlines`` logger to emit DEBUG messages to stderr."""
    flowlines_logger = logging.getLogger("flowlines")
    if flowlines_logger.level <= logging.DEBUG and flowlines_logger.handlers:
        return
    flowlines_logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler()
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(logging.Formatter("[%(name)s] %(message)s"))
    flowlines_logger.addHandler(handler)


class Flowlines:
    """Flowlines SDK singleton for LLM observability.

    Only one instance may be created. Attempting to instantiate a second
    one raises a ``RuntimeError``.

    The constructor initializes the SDK based on the selected mode:

    - **Mode A** (default): No existing OpenTelemetry setup. Creates a
      ``TracerProvider``, adds a span processor, sets it as global, and
      registers all available instrumentors.
    - **Mode B1** (``has_external_otel=True``): Skips ``TracerProvider``
      creation and instrumentor registration. The user is responsible for
      calling ``create_span_processor()`` and ``get_instrumentors()`` to
      integrate with their own setup.
    - **Mode B2** (``has_traceloop=True``): Assumes Traceloop is already
      initialized. Adds a span processor to the existing ``TracerProvider``.
      Skips instrumentor registration.

    Args:
        api_key: The Flowlines API key.
        endpoint: The Flowlines backend URL. Must use HTTPS unless targeting
            localhost/127.0.0.1.
        has_external_otel: If True, skips ``TracerProvider`` creation and
            instrumentor registration (Mode B1).
        has_traceloop: If True, assumes Traceloop is already initialized
            and only adds the Flowlines span processor (Mode B2).
        verbose: If True, enables DEBUG-level logging to stderr for the
            ``flowlines`` logger. Useful for troubleshooting missing spans.

    Raises:
        RuntimeError: If a second instance is created, or if
            ``has_traceloop=True`` but no SDK ``TracerProvider`` is found.
        ValueError: If the endpoint is not HTTPS and not localhost/127.0.0.1,
            or if both ``has_external_otel`` and ``has_traceloop`` are True.
    """

    _instance: Flowlines | None = None
    _lock: threading.RLock = threading.RLock()

    def __init__(
        self,
        api_key: str,
        endpoint: str = _DEFAULT_ENDPOINT,
        has_external_otel: bool = False,
        has_traceloop: bool = False,
        verbose: bool = False,
    ) -> None:
        with Flowlines._lock:
            if Flowlines._instance is not None:
                msg = "Flowlines is a singleton — only one instance may be created."
                raise RuntimeError(msg)

            if not api_key or not api_key.strip():
                msg = "api_key must not be empty."
                raise ValueError(msg)

            if has_external_otel and has_traceloop:
                msg = "has_external_otel and has_traceloop are mutually exclusive."
                raise ValueError(msg)

            _validate_endpoint(endpoint)

            if verbose:
                _configure_verbose_logging()

            _apply_patches()

            self._api_key = api_key
            self._endpoint = endpoint
            self._has_external_otel = has_external_otel
            self._has_traceloop = has_traceloop
            self._processor: BatchSpanProcessor | None = None
            self._provider: TracerProvider | None = None
            self._shutdown_called = False

            Flowlines._instance = self

            logger.debug("Initializing Flowlines SDK (endpoint=%s)", endpoint)

            try:
                if has_external_otel:
                    # Mode B1: user manages TracerProvider and instrumentors.
                    logger.debug(
                        "Mode B1: external OpenTelemetry setup; "
                        "skipping TracerProvider and instrumentor registration"
                    )
                elif has_traceloop:
                    # Mode B2: Traceloop is already initialized.
                    logger.debug(
                        "Mode B2: Traceloop detected; "
                        "adding span processor to existing TracerProvider"
                    )
                    provider = trace.get_tracer_provider()
                    if not isinstance(provider, TracerProvider):
                        msg = (
                            "has_traceloop=True but no SDK TracerProvider found. "
                            "Make sure Traceloop is initialized before Flowlines."
                        )
                        raise RuntimeError(msg)
                    processor = self.create_span_processor()
                    provider.add_span_processor(processor)
                else:
                    # Mode A: no existing OTEL setup.
                    logger.debug(
                        "Mode A: creating TracerProvider and registering instrumentors"
                    )
                    self._provider = TracerProvider()
                    processor = self.create_span_processor()
                    self._provider.add_span_processor(processor)
                    trace.set_tracer_provider(self._provider)

                    for instrumentor in self.get_instrumentors():
                        instrumentor.instrument(tracer_provider=self._provider)
            except Exception:
                Flowlines._instance = None
                raise

            logger.debug("Flowlines SDK initialized successfully")
            atexit.register(self.shutdown)

    def create_span_processor(self) -> SpanProcessor:
        """Create a span processor wrapping the ``FlowlinesExporter``.

        Returns a ``FlowlinesSpanProcessor`` that injects context attributes
        (``flowlines.user_id``, ``flowlines.session_id``,
        ``flowlines.agent_id``) and delegates to a ``BatchSpanProcessor``
        for export.

        May only be called once. In Mode A and Mode B2 it is called
        automatically during ``__init__``.  In Mode B1, the user should
        call it exactly once.

        Returns:
            A configured ``SpanProcessor``.

        Raises:
            RuntimeError: If called more than once.
        """
        with Flowlines._lock:
            if self._processor is not None:
                msg = "create_span_processor() has already been called."
                raise RuntimeError(msg)
            exporter = FlowlinesExporter(api_key=self._api_key, endpoint=self._endpoint)
            batch_processor = BatchSpanProcessor(exporter)
            processor = FlowlinesSpanProcessor(batch_processor)
            self._processor = batch_processor
            return processor

    def get_instrumentors(self) -> list[Any]:
        """Return instances of all available OpenLLMetry instrumentors.

        Each instrumentor is only included if its underlying library is
        installed. For example, ``OpenAIInstrumentor`` is only returned if
        the ``openai`` package is importable.

        Returns:
            A list of instrumentor instances.
        """
        instrumentors: list[Any] = []
        for lib_spec, module_path, class_name in _INSTRUMENTOR_REGISTRY:
            if importlib.util.find_spec(lib_spec) is None:
                logger.debug(
                    "Instrumentor skipped: %s (library not installed)",
                    lib_spec,
                )
                continue
            try:
                instrumentors.append(_load_instrumentor(module_path, class_name))
                logger.debug("Instrumentor loaded: %s", class_name)
            except (ImportError, AttributeError) as exc:
                logger.debug("Instrumentor failed to load: %s (%s)", class_name, exc)
                continue
        logger.debug("Total instrumentors loaded: %d", len(instrumentors))
        return instrumentors

    def shutdown(self) -> None:
        """Flush pending spans and shut down the processor/provider.

        In Modes B1/B2, calls ``force_flush()`` then ``shutdown()`` on the
        span processor. In Mode A, calls them on the ``TracerProvider``.

        Safe to call multiple times (idempotent). Also registered as an
        ``atexit`` handler automatically.
        """
        if self._shutdown_called:
            return
        self._shutdown_called = True

        logger.debug("Shutting down Flowlines SDK")
        if self._has_external_otel or self._has_traceloop:
            # Modes B1/B2: flush and shut down the span processor.
            if self._processor is not None:
                self._processor.force_flush()
                self._processor.shutdown()  # type: ignore[no-untyped-call]
        else:
            # Mode A: flush and shut down the TracerProvider.
            if self._provider is not None:
                self._provider.force_flush()
                self._provider.shutdown()
        logger.debug("Flowlines SDK shut down complete")

    @contextmanager
    def context(
        self,
        *,
        user_id: str,
        session_id: str | None = None,
        agent_id: str | None = None,
    ) -> Generator[None, None, None]:
        """Context manager that sets Flowlines attributes on LLM spans.

        Any spans created within the block will have the specified
        ``flowlines.user_id`` and optionally ``flowlines.session_id``
        and ``flowlines.agent_id`` attributes. Nesting uses replace
        semantics — an inner context fully replaces the outer one.

        Args:
            user_id: The user ID to attach to spans (required).
            session_id: The session ID to attach to spans.
            agent_id: The agent ID to attach to spans.
        """
        token = _set_context(user_id=user_id, session_id=session_id, agent_id=agent_id)
        try:
            yield
        finally:
            _clear_context(token)

    @staticmethod
    def set_context(
        *,
        user_id: str | None = None,
        session_id: str | None = None,
        agent_id: str | None = None,
    ) -> Token[dict[str, str] | None]:
        """Set Flowlines context attributes imperatively.

        Returns a token that must be passed to ``clear_context`` to restore
        the previous state.

        Args:
            user_id: The user ID to attach to spans.
            session_id: The session ID to attach to spans.
            agent_id: The agent ID to attach to spans.

        Returns:
            A ``Token`` for restoring the previous context.
        """
        return _set_context(user_id=user_id, session_id=session_id, agent_id=agent_id)

    @staticmethod
    def clear_context(token: Token[dict[str, str] | None]) -> None:
        """Restore the previous Flowlines context.

        Args:
            token: The token returned by ``set_context``.
        """
        _clear_context(token)

    @classmethod
    def _reset(cls) -> None:
        """Reset the singleton. For testing only."""
        with cls._lock:
            instance = cls._instance
            if instance is not None:
                instance.shutdown()
                atexit.unregister(instance.shutdown)
            cls._instance = None
