from __future__ import annotations

import os
import logging
from typing import Any, TYPE_CHECKING, Literal

from opentelemetry.instrumentation.threading import ThreadingInstrumentor
from opentelemetry.propagate import set_global_textmap
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from .otel import setup_tracing, TracingRuntime
from .context import enable_context_processor, merge_context_properties
from .registry import auto_instrument_all, instrument_httpx

logger = logging.getLogger("impact.sdk")

if TYPE_CHECKING:
    from typing import Protocol
    
    class ASGIApp(Protocol):
        async def __call__(self, scope, receive, send): ...


InitMode = Literal["auto", "bootstrap", "attach"]

_DEFAULT_DOMAIN = "beta.impact.ai"

_initialized = False
_runtime: TracingRuntime | None = None


def init(
    *,
    api_key: str | None = None,
    api_endpoint: str | None = None,
    mode: InitMode = "auto",
    disable_batch: bool = False,
    disable_logs: bool = False,
) -> None:
    """Initialize OpenTelemetry tracing and auto-instrument GenAI SDKs.
    Args:
        api_key: API key for authentication. Falls back to IMPACT_API_KEY env var.
        api_endpoint: Impact base API endpoint. Falls back to IMPACT_BASE_URL env var.
        mode: "auto" (default), "bootstrap", or "attach".
            - auto: attach to existing TracerProvider (eg: Langfuse, Datadog, etc.), else bootstrap.
            - bootstrap: always create and set a new TracerProvider / LoggerProvider.
            - attach: require an existing TracerProvider; do not replace it.
        disable_batch: If True, use SimpleSpanProcessor / SimpleLogRecordProcessor.
        disable_logs: If True, skip configuring logs (disables GenAI message content capture).

    Raises:
        ValueError: If api_endpoint is missing, or if mode="attach" and no TracerProvider exists.
    """
    global _initialized, _runtime

    if _initialized:
        return

    # Set default environment variables for latest GenAI semantics + capture message content.
    os.environ.setdefault("OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental")
    os.environ.setdefault("OTEL_INSTRUMENTATION_OPENAI_AGENTS_CAPTURE_CONTENT", "true")
    os.environ.setdefault("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "SPAN_AND_EVENT")
    os.environ.setdefault("AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED", "true")    # Microsoft Foundry
    os.environ.setdefault("ENABLE_INSTRUMENTATION", "true")    # Microsoft Agent Framework
    os.environ.setdefault("ENABLE_SENSITIVE_DATA", "true")    # Microsoft Agent Framework
    os.environ.setdefault("HAYSTACK_CONTENT_TRACING_ENABLED", "true")    # Haystack

    # Resolve API key
    api_key = api_key or os.getenv("IMPACT_API_KEY")
    if not api_key:
        raise ValueError("api_key is required — pass to init() or set IMPACT_API_KEY")

    # Resolve Endpoint: (1) explicit override, (2) env var (IMPACT_BASE_URL), or (3) derived from API key
    api_endpoint = api_endpoint or os.getenv("IMPACT_BASE_URL")
    if not api_endpoint:
        parts = api_key.split("_", 3)
        if len(parts) >= 3 and parts[0] == "impact":
            domain = os.getenv("IMPACT_DOMAIN", _DEFAULT_DOMAIN)
            api_endpoint = f"https://api.{parts[1]}.{domain}"
        else:
            raise ValueError("Could not resolve endpoint from API key")

    # Programmatic headers
    headers: dict[str, str] = {}
    if api_key:
        headers["authorization"] = f"Bearer {api_key}"

    # Setup tracing/logs (bootstrap or attach)
    _runtime = setup_tracing(
        api_endpoint=api_endpoint,
        headers=headers,
        mode=mode,
        disable_batch=disable_batch,
        disable_logs=disable_logs,
    )

    # Only set propagators when we bootstrap (do not override user's propagation in attach mode).
    if not _runtime.attached:
        set_global_textmap(TraceContextTextMapPropagator())

    # Enable Impact context propagation
    enable_context_processor()

    # Enable thread context propagation
    ThreadingInstrumentor().instrument()

    # Auto-instrument HTTPX HTTP client for outbound requests
    instrument_httpx()

    # Auto-instrument supported GenAI SDKs (best-effort, non-fatal on conflicts).
    auto_instrument_all()

    _initialized = True


def context(*, user_id: str | None = None, interaction_id: str | None = None, version_id: str | None = None, **extras: Any) -> None:
    """Attach context for all spans in current execution context."""
    props: dict[str, Any] = {}
    if user_id is not None:
        props["user_id"] = user_id
    if interaction_id is not None:
        props["interaction_id"] = interaction_id
    if version_id is not None:
        props["version_id"] = version_id
    if extras:
        props.update(extras)

    if props:
        merge_context_properties(props)


def instrument_asgi_app(app: Any) -> Any:
    """Wrap an ASGI application with OpenTelemetry middleware."""
    try:
        from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
        return OpenTelemetryMiddleware(app)
    except ImportError:
        # If ASGI instrumentation not available, return app unchanged
        return app

def shutdown() -> None:
    """Flush and shutdown Impact's telemetry components, will NOT shutdown the user's TracerProvider / LoggerProvider."""
    global _runtime, _initialized

    rt = _runtime
    _runtime = None
    _initialized = False

    if rt is None:
        return

    try:
        rt.shutdown()
    except Exception:
        logger.exception("Failed to shutdown Impact telemetry runtime")
