"""Optional observability/instrumentation helpers.

This module is safe to import without optional dependencies installed.
Call :func:`instrument_ketju` when you explicitly want instrumentation.
"""

from __future__ import annotations

from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
import logging
from typing import Literal, Mapping, Optional

AttributeValue = str | bool | int | float | list[str] | list[bool] | list[int] | list[float]

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class InstrumentationResult:
    """What was enabled successfully."""

    logfire_configured: bool = False
    pydantic_ai_instrumented: bool = False
    llamaindex_instrumented: bool = False


@contextmanager
def start_span(
    name: str,
    *,
    attributes: Mapping[str, AttributeValue] | None = None,
):
    """Start a span if OpenTelemetry is available.

    This uses the **application's configured** global tracer provider. If the host
    app has instrumented lower-level libraries (e.g. OpenAI/httpx), those spans
    will naturally show up as children of this span when they run inside it.
    """

    try:
        from opentelemetry import trace  # type: ignore
    except ModuleNotFoundError:
        logger.debug("OpenTelemetry not installed; start_span is a no-op")
        with nullcontext():
            yield None
        return

    tracer = trace.get_tracer("ketju")
    logger.debug("Using tracer provider for ketju spans: %r", tracer)
    with tracer.start_as_current_span(name) as span:
        if attributes:
            for k, v in attributes.items():
                try:
                    span.set_attribute(k, v)
                except Exception:
                    # Never let observability break the business logic.
                    pass
        yield span


def add_span_event(
    name: str,
    *,
    attributes: Mapping[str, AttributeValue] | None = None,
) -> None:
    """Add an event to the current span (no-op if none)."""

    try:
        from opentelemetry import trace  # type: ignore
    except ModuleNotFoundError:
        return

    try:
        span = trace.get_current_span()
        # Avoid adding events to the no-op/invalid span.
        if not span.get_span_context().is_valid:  # type: ignore
            return
        span.add_event(name, attributes=dict(attributes or {}))
    except Exception:
        # Observability must never break the business logic.
        return


def configure_otlp_tracing(
    *,
    service_name: str = "ketju",
    protocol: Literal["grpc", "http/protobuf"] = "http/protobuf",
    otlp_endpoint: str,
    otlp_headers: Optional[Mapping[str, str]] = None,
    strict: bool = False,
) -> bool:
    """Application helper: configure OTLP tracing exporter.

    Libraries typically should **not** configure exporters/providers automatically.
    This helper is provided for apps/examples that want a batteries-included setup.

    Install deps with:
      `uv sync --extra observability`
    """

    try:
        from opentelemetry import trace  # type: ignore
        from opentelemetry.sdk.resources import Resource  # type: ignore
        from opentelemetry.sdk.trace import TracerProvider  # type: ignore
        from opentelemetry.sdk.trace.export import BatchSpanProcessor  # type: ignore
        if protocol == "grpc":
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (  # type: ignore
                OTLPSpanExporter,
            )
            # gRPC endpoint should be base host:port (no /v1/traces).
            if "/v1/traces" in otlp_endpoint:
                otlp_endpoint = otlp_endpoint.split("/v1/traces")[0]
        else:
            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (  # type: ignore
                OTLPSpanExporter,
            )
            # HTTP endpoint should include /v1/traces for most collectors.
            if "/v1/traces" not in otlp_endpoint.rstrip("/"):
                otlp_endpoint = otlp_endpoint.rstrip("/") + "/v1/traces"

        headers = dict(otlp_headers or {})
        provider = TracerProvider(resource=Resource.create({"service.name": service_name}))
        exporter = OTLPSpanExporter(endpoint=otlp_endpoint, headers=headers or None)
        provider.add_span_processor(BatchSpanProcessor(exporter))
        trace.set_tracer_provider(provider)
        logger.info("Configured OTLP tracing (%s) to %s", protocol, otlp_endpoint)
        return True
    except ModuleNotFoundError:
        if strict:
            raise ModuleNotFoundError(
                "OpenTelemetry SDK/exporter not installed. Install with: `uv sync --extra observability`"
            )
        return False
    except Exception:
        if strict:
            raise
        return False


def instrument_ketju(
    *,
    service_name: str = "ketju",
    enable_logfire: bool = True,
    instrument_pydantic_ai: bool = True,
    instrument_llamaindex: bool = True,
    quiet: bool = True,
    strict: bool = False,
) -> InstrumentationResult:
    """Enable Ketju instrumentation hooks.

    This intentionally **does not configure exporters/providers**; it only enables
    instrumentation (and optionally Logfire configuration) if the dependencies exist.
    """

    if quiet:
        for name in ("opentelemetry", "opentelemetry.instrumentation", "httpx"):
            logging.getLogger(name).setLevel(logging.WARNING)

    result = InstrumentationResult()

    # Logfire (optional)
    try:
        import logfire  # type: ignore

        if enable_logfire:
            logfire.configure(service_name=service_name)
            result = InstrumentationResult(
                logfire_configured=True,
                pydantic_ai_instrumented=False,
                llamaindex_instrumented=False,
            )

            if instrument_pydantic_ai:
                try:
                    logfire.instrument_pydantic_ai()
                    result = InstrumentationResult(
                        logfire_configured=True,
                        pydantic_ai_instrumented=True,
                        llamaindex_instrumented=result.llamaindex_instrumented,
                    )
                except Exception:
                    if strict:
                        raise
    except ModuleNotFoundError:
        if strict and enable_logfire:
            raise ModuleNotFoundError(
                "logfire is not installed. Install with: `uv sync --extra observability`"
            )

    # OpenTelemetry llamaindex instrumentation (optional)
    if instrument_llamaindex:
        try:
            from opentelemetry.instrumentation.llamaindex import (  # type: ignore
                LlamaIndexInstrumentor,
            )

            LlamaIndexInstrumentor().instrument()
            result = InstrumentationResult(
                logfire_configured=result.logfire_configured,
                pydantic_ai_instrumented=result.pydantic_ai_instrumented,
                llamaindex_instrumented=True,
            )
        except ModuleNotFoundError:
            if strict:
                raise ModuleNotFoundError(
                    "OpenTelemetry llamaindex instrumentation is not installed. "
                    "Install with: `uv sync --extra observability`"
                )
        except Exception:
            if strict:
                raise

    return result

