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

# Always emit ERROR+ messages to stderr, even when verbose mode is off.
if not any(getattr(h, "_flowlines_error_handler", False) for h in logger.handlers):
    _error_handler = logging.StreamHandler()
    _error_handler.setLevel(logging.ERROR)
    _error_handler.setFormatter(logging.Formatter("[%(name)s] %(message)s"))
    _error_handler._flowlines_error_handler = True  # type: ignore[attr-defined]
    logger.addHandler(_error_handler)

_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 _log_memory_error(exc: Exception) -> None:
    """Log a memory error at the appropriate level.

    HTTP 5xx, 429 and network errors are logged at ERROR level (always
    visible).  Other HTTP errors (4xx) are logged at WARNING level
    (visible only in verbose mode).
    """
    from flowlines._memory import FlowlinesMemoryError

    if isinstance(exc, FlowlinesMemoryError):
        sc = exc.status_code
        if sc is None or sc == 429 or sc >= 500:
            logger.error("%s", exc)
        else:
            logger.warning("%s", exc)
    else:
        logger.error("%s", exc)


def _log_session_error(exc: Exception) -> None:
    """Log a session error at the appropriate level.

    HTTP 5xx, 429 and network errors are logged at ERROR level (always
    visible).  Other HTTP errors (4xx) are logged at WARNING level
    (visible only in verbose mode).
    """
    from flowlines._session import FlowlinesSessionError

    if isinstance(exc, FlowlinesSessionError):
        sc = exc.status_code
        if sc is None or sc == 429 or sc >= 500:
            logger.error("%s", exc)
        else:
            logger.warning("%s", exc)
    else:
        logger.error("%s", exc)


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:
        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")
        try:
            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()
        except Exception:
            logger.error("Error during Flowlines SDK shutdown", exc_info=True)
            return
        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).

        Errors are logged but never propagated so that callers are not
        disrupted by transient backend or network issues.

        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 or an error occurred.
        """
        from flowlines._memory import FlowlinesMemoryError, fetch_memory

        try:
            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,
            )
        except FlowlinesMemoryError as exc:
            _log_memory_error(exc)
            return None

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

        Errors are logged but never propagated so that callers are not
        disrupted by transient backend or network issues.

        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 or an error occurred.
        """
        from flowlines._memory import FlowlinesMemoryError, afetch_memory

        try:
            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,
            )
        except FlowlinesMemoryError as exc:
            _log_memory_error(exc)
            return None

    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.  Errors are logged
        but never propagated so that callers are not disrupted by
        transient backend or network issues.

        Args:
            user_id: The user ID (required).
            session_id: The session ID (required).
        """
        from flowlines._session import FlowlinesSessionError
        from flowlines._session import end_session as _end_session

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

        try:
            _end_session(
                api_key=self._api_key,
                api_endpoint=self._api_endpoint,
                user_id=user_id,
                session_id=session_id,
            )
        except FlowlinesSessionError as exc:
            _log_session_error(exc)

    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.  Errors are logged
        but never propagated so that callers are not disrupted by
        transient backend or network issues.

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

        from flowlines._session import FlowlinesSessionError
        from flowlines._session import aend_session as _aend_session

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

        try:
            await _aend_session(
                api_key=self._api_key,
                api_endpoint=self._api_endpoint,
                user_id=user_id,
                session_id=session_id,
            )
        except FlowlinesSessionError as exc:
            _log_session_error(exc)


# ---------------------------------------------------------------------------
# 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).

    Errors from the backend or network are logged but never propagated,
    so calling code is not disrupted by transient issues.

    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 or an error occurred.

    Raises:
        RuntimeError: If Flowlines has not been initialized.
    """
    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).

    Errors from the backend or network are logged but never propagated,
    so calling code is not disrupted by transient issues.

    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 or an error occurred.

    Raises:
        RuntimeError: If Flowlines has not been initialized.
    """
    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).

    Errors from the backend or network are logged but never propagated,
    so calling code is not disrupted by transient issues.

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

    Raises:
        RuntimeError: If Flowlines has not been initialized.
    """
    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).

    Errors from the backend or network are logged but never propagated,
    so calling code is not disrupted by transient issues.

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

    Raises:
        RuntimeError: If Flowlines has not been initialized.
    """
    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
