"""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_INGEST_ENDPOINT = "https://ingest.flowlines.ai"
_DEFAULT_API_ENDPOINT = "https://api.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"),
    ("groq", "opentelemetry.instrumentation.groq", "GroqInstrumentor"),
    ("mistralai", "opentelemetry.instrumentation.mistralai", "MistralAiInstrumentor"),
    ("ollama", "opentelemetry.instrumentation.ollama", "OllamaInstrumentor"),
    ("replicate", "opentelemetry.instrumentation.replicate", "ReplicateInstrumentor"),
    (
        "transformers",
        "opentelemetry.instrumentation.transformers",
        "TransformersInstrumentor",
    ),
    ("boto3", "opentelemetry.instrumentation.sagemaker", "SageMakerInstrumentor"),
    ("crewai", "opentelemetry.instrumentation.crewai", "CrewAIInstrumentor"),
    ("agno", "opentelemetry.instrumentation.agno", "AgnoInstrumentor"),
    ("haystack", "opentelemetry.instrumentation.haystack", "HaystackInstrumentor"),
    ("lancedb", "opentelemetry.instrumentation.lancedb", "LanceInstrumentor"),
    ("marqo", "opentelemetry.instrumentation.marqo", "MarqoInstrumentor"),
    ("pymilvus", "opentelemetry.instrumentation.milvus", "MilvusInstrumentor"),
    ("weaviate", "opentelemetry.instrumentation.weaviate", "WeaviateInstrumentor"),
    (
        "aleph_alpha_client",
        "opentelemetry.instrumentation.alephalpha",
        "AlephAlphaInstrumentor",
    ),
    ("voyageai", "opentelemetry.instrumentation.voyageai", "VoyageAIInstrumentor"),
    (
        "ibm_watson_machine_learning",
        "opentelemetry.instrumentation.watsonx",
        "WatsonxInstrumentor",
    ),
    ("writer", "opentelemetry.instrumentation.writer", "WriterInstrumentor"),
    (
        "agents",
        "opentelemetry.instrumentation.openai_agents",
        "OpenAIAgentsInstrumentor",
    ),
]


_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)


# ---------------------------------------------------------------------------
# Private implementation class
# ---------------------------------------------------------------------------


class _Flowlines:
    """Private implementation of the Flowlines SDK.

    Not intended for direct use. Use the module-level ``init()`` function.
    """

    def __init__(
        self,
        api_key: str,
        has_external_otel: bool = False,
        has_traceloop: bool = False,
        verbose: bool = False,
        ingest_endpoint: str = _DEFAULT_INGEST_ENDPOINT,
        api_endpoint: str = _DEFAULT_API_ENDPOINT,
    ) -> None:
        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(ingest_endpoint)
        _validate_endpoint(api_endpoint)

        if verbose:
            _configure_verbose_logging()

        _apply_patches()

        self._api_key = api_key
        self._ingest_endpoint = ingest_endpoint
        self._api_endpoint = api_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

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

        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 get_instrumentors():
                instrumentor.instrument(tracer_provider=self._provider)

        logger.debug("Flowlines SDK initialized successfully")

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

        May only be called once. In Mode A and Mode B2 it is called
        automatically during ``init()``. In Mode B1, the user should
        call the module-level ``create_span_processor()`` exactly once.

        Returns:
            A configured ``SpanProcessor``.

        Raises:
            RuntimeError: If called more than once.
        """
        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._ingest_endpoint
        )
        batch_processor = BatchSpanProcessor(exporter)
        processor = FlowlinesSpanProcessor(batch_processor)
        self._processor = batch_processor
        return processor

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

        Safe to call multiple times (idempotent).
        """
        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")

    def get_memory(
        self,
        user_id: str,
        *,
        session_id: str | None = None,
        agent_id: str | None = None,
        view: str | None = None,
    ) -> str | None:
        """Retrieve memory from the Flowlines backend (synchronous).

        Args:
            user_id: The user ID (required).
            session_id: Optional session ID filter.
            agent_id: Optional agent ID filter.
            view: Optional view name.

        Returns:
            A JSON string of the ``memory`` object from the response, or
            ``None`` if the server returned HTTP 404.

        Raises:
            FlowlinesMemoryError: On HTTP errors (other than 404), network
                failures, or malformed responses.
        """
        from flowlines._memory import fetch_memory

        return fetch_memory(
            api_key=self._api_key,
            api_endpoint=self._api_endpoint,
            user_id=user_id,
            session_id=session_id,
            agent_id=agent_id,
            view=view,
        )

    async def aget_memory(
        self,
        user_id: str,
        *,
        session_id: str | None = None,
        agent_id: str | None = None,
        view: str | None = None,
    ) -> str | None:
        """Retrieve memory from the Flowlines backend (asynchronous).

        Args:
            user_id: The user ID (required).
            session_id: Optional session ID filter.
            agent_id: Optional agent ID filter.
            view: Optional view name.

        Returns:
            A JSON string of the ``memory`` object from the response, or
            ``None`` if the server returned HTTP 404.

        Raises:
            FlowlinesMemoryError: On HTTP errors (other than 404), network
                failures, or malformed responses.
        """
        from flowlines._memory import afetch_memory

        return await afetch_memory(
            api_key=self._api_key,
            api_endpoint=self._api_endpoint,
            user_id=user_id,
            session_id=session_id,
            agent_id=agent_id,
            view=view,
        )

    def end_session(self, user_id: str, *, session_id: str) -> None:
        """Mark a session as ended on the Flowlines backend (synchronous).

        Flushes pending spans before notifying the backend so that all
        traces for the session are delivered first.

        Args:
            user_id: The user ID (required).
            session_id: The session ID (required).

        Raises:
            FlowlinesSessionError: On HTTP errors (other than 404), network
                failures, or malformed responses.
        """
        from flowlines._session import end_session as _end_session

        if self._processor is not None:
            self._processor.force_flush()

        _end_session(
            api_key=self._api_key,
            api_endpoint=self._api_endpoint,
            user_id=user_id,
            session_id=session_id,
        )

    async def aend_session(self, user_id: str, *, session_id: str) -> None:
        """Mark a session as ended on the Flowlines backend (asynchronous).

        Flushes pending spans before notifying the backend so that all
        traces for the session are delivered first.

        Args:
            user_id: The user ID (required).
            session_id: The session ID (required).

        Raises:
            FlowlinesSessionError: On HTTP errors (other than 404), network
                failures, or malformed responses.
        """
        import asyncio

        from flowlines._session import aend_session as _aend_session

        if self._processor is not None:
            await asyncio.to_thread(self._processor.force_flush)

        await _aend_session(
            api_key=self._api_key,
            api_endpoint=self._api_endpoint,
            user_id=user_id,
            session_id=session_id,
        )


# ---------------------------------------------------------------------------
# Module-level singleton state
# ---------------------------------------------------------------------------

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


# ---------------------------------------------------------------------------
# Public module-level functions
# ---------------------------------------------------------------------------


def init(
    api_key: str,
    has_external_otel: bool = False,
    has_traceloop: bool = False,
    verbose: bool = False,
    ingest_endpoint: str = _DEFAULT_INGEST_ENDPOINT,
    api_endpoint: str = _DEFAULT_API_ENDPOINT,
) -> None:
    """Initialize the Flowlines SDK.

    Safe to call multiple times — subsequent calls are silent no-ops.

    The SDK operates in one of three modes:

    - **Mode A** (default): Creates a ``TracerProvider``, registers all
      available instrumentors automatically.
    - **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``.

    Args:
        api_key: The Flowlines API key.
        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.
        ingest_endpoint: The Flowlines ingest backend URL. Must use HTTPS
            unless targeting localhost/127.0.0.1.
        api_endpoint: The Flowlines API endpoint for memory retrieval.
            Must use HTTPS unless targeting localhost/127.0.0.1.

    Raises:
        RuntimeError: If ``has_traceloop=True`` but no SDK
            ``TracerProvider`` is found.
        ValueError: If an endpoint is not HTTPS and not localhost/127.0.0.1,
            or if both ``has_external_otel`` and ``has_traceloop`` are True.
    """
    global _instance  # noqa: PLW0603
    with _lock:
        if _instance is not None:
            logger.debug("Flowlines is already initialized — skipping")
            return

        try:
            _instance = _Flowlines(
                api_key=api_key,
                ingest_endpoint=ingest_endpoint,
                has_external_otel=has_external_otel,
                has_traceloop=has_traceloop,
                verbose=verbose,
                api_endpoint=api_endpoint,
            )
        except Exception:
            _instance = None
            raise

        atexit.register(shutdown)


def shutdown() -> None:
    """Flush pending spans and shut down the SDK.

    No-op if Flowlines has not been initialized. Safe to call multiple
    times (idempotent). Registered as an ``atexit`` handler automatically.
    """
    if _instance is None:
        return
    _instance.shutdown()


def create_span_processor() -> 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 Flowlines has not been initialized, or if
            called more than once.
    """
    with _lock:
        if _instance is None:
            msg = (
                "Flowlines has not been initialized."
                " Call flowlines.init(api_key=...) first."
            )
            raise RuntimeError(msg)
        return _instance.create_span_processor()


def get_memory(
    user_id: str,
    *,
    session_id: str | None = None,
    agent_id: str | None = None,
    view: str | None = None,
) -> str | None:
    """Retrieve memory from the Flowlines backend (synchronous).

    Args:
        user_id: The user ID (required).
        session_id: Optional session ID filter.
        agent_id: Optional agent ID filter.
        view: Optional view name.

    Returns:
        A JSON string of the ``memory`` object from the response, or
        ``None`` if the server returned HTTP 404.

    Raises:
        RuntimeError: If Flowlines has not been initialized.
        FlowlinesMemoryError: On HTTP errors (other than 404), network
            failures, or malformed responses.
    """
    if _instance is None:
        msg = (
            "Flowlines has not been initialized."
            " Call flowlines.init(api_key=...) first."
        )
        raise RuntimeError(msg)
    return _instance.get_memory(
        user_id, session_id=session_id, agent_id=agent_id, view=view
    )


async def aget_memory(
    user_id: str,
    *,
    session_id: str | None = None,
    agent_id: str | None = None,
    view: str | None = None,
) -> str | None:
    """Retrieve memory from the Flowlines backend (asynchronous).

    Args:
        user_id: The user ID (required).
        session_id: Optional session ID filter.
        agent_id: Optional agent ID filter.
        view: Optional view name.

    Returns:
        A JSON string of the ``memory`` object from the response, or
        ``None`` if the server returned HTTP 404.

    Raises:
        RuntimeError: If Flowlines has not been initialized.
        FlowlinesMemoryError: On HTTP errors (other than 404), network
            failures, or malformed responses.
    """
    if _instance is None:
        msg = (
            "Flowlines has not been initialized."
            " Call flowlines.init(api_key=...) first."
        )
        raise RuntimeError(msg)
    return await _instance.aget_memory(
        user_id, session_id=session_id, agent_id=agent_id, view=view
    )


def end_session(user_id: str, *, session_id: str) -> None:
    """Mark a session as ended on the Flowlines backend (synchronous).

    Args:
        user_id: The user ID (required).
        session_id: The session ID (required).

    Raises:
        RuntimeError: If Flowlines has not been initialized.
        FlowlinesSessionError: On HTTP errors (other than 404), network
            failures, or malformed responses.
    """
    if _instance is None:
        msg = (
            "Flowlines has not been initialized."
            " Call flowlines.init(api_key=...) first."
        )
        raise RuntimeError(msg)
    _instance.end_session(user_id, session_id=session_id)


async def aend_session(user_id: str, *, session_id: str) -> None:
    """Mark a session as ended on the Flowlines backend (asynchronous).

    Args:
        user_id: The user ID (required).
        session_id: The session ID (required).

    Raises:
        RuntimeError: If Flowlines has not been initialized.
        FlowlinesSessionError: On HTTP errors (other than 404), network
            failures, or malformed responses.
    """
    if _instance is None:
        msg = (
            "Flowlines has not been initialized."
            " Call flowlines.init(api_key=...) first."
        )
        raise RuntimeError(msg)
    await _instance.aend_session(user_id, session_id=session_id)


def get_instrumentors() -> 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


@contextmanager
def context(
    *,
    user_id: str,
    session_id: str,
    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``, ``flowlines.session_id`` and optionally
    ``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 (required).
        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)


def set_context(
    *,
    user_id: str,
    session_id: str,
    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 (required).
        session_id: The session ID to attach to spans (required).
        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,
    )


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)


def _reset() -> None:  # pyright: ignore[reportUnusedFunction]
    """Reset the singleton. For testing only."""
    global _instance  # noqa: PLW0603
    with _lock:
        if _instance is not None:
            _instance.shutdown()
            atexit.unregister(shutdown)
        _instance = None
