from __future__ import annotations

import logging
import threading
from typing import TYPE_CHECKING, Optional

import grpc
from observability_common.auth_hooks import AuthHook, NoAuthHook

if TYPE_CHECKING:
    import networkx as nwx
from observability_common.proto import observability_pb2_grpc as observability_grpc
from observability_common.proto.observability_pb2 import (
    Computation,
    Computations,
    Dependencies,
    Edge,
    Event,
    Events,
    SliceRequest,
    SpanRequest,
    TraceGraph,
    TraceRequest,
)
from observability_common.tls import TLSConfigWithFallback
from pydantic import Field
from pydantic_settings import SettingsConfigDict

from .utils import to_digraph

logger = logging.getLogger(__name__)

# Cached channel for observability API (gRPC channels are thread-safe)
_channel_lock = threading.Lock()
_cached_channel: Optional[grpc.Channel] = None


class ObservabilityConfig(TLSConfigWithFallback):
    """
    Global configuration for observability API.

    Inherits TLS settings from TLSConfigWithFallback, which falls back to
    shared TLS_ env vars (TLS_CA_PATH, TLS_CERT_PATH, TLS_KEY_PATH).

    Environment variables:
        OBSERVABILITY_URL=localhost:10085
    """

    model_config = SettingsConfigDict(
        env_prefix="OBSERVABILITY_",
        env_file=".env",
        extra="ignore",
        arbitrary_types_allowed=True,
    )

    url: str = Field(
        default="localhost:10085",
        description="Observability server URL (host:port)",
    )
    auth_hook: AuthHook = Field(
        default_factory=NoAuthHook,
        description="Auth hook for JWT/OIDC authentication (same as reference_monitor)",
    )

    def create_channel(self) -> grpc.Channel:
        """Create a gRPC channel with configured TLS settings."""
        ca_cert = self.load_ca_cert()
        client_cert = self.load_cert()
        client_key = self.load_key()

        logger.debug(f"TLS config: ca={self.ca_path} ({ca_cert is not None}), "
                    f"cert={self.cert_path} ({client_cert is not None}), "
                    f"key={self.key_path} ({client_key is not None})")

        if client_cert and client_key and ca_cert:
            logger.debug(f"Creating mTLS channel to {self.url}")
            # mTLS: client cert + server verification
            credentials = grpc.ssl_channel_credentials(
                root_certificates=ca_cert,
                private_key=client_key,
                certificate_chain=client_cert,
            )
            return grpc.secure_channel(self.url, credentials)
        elif ca_cert:
            logger.warning(f"Creating TLS channel to {self.url} (no client cert - mTLS auth won't work)")
            # TLS only: just verify server
            credentials = grpc.ssl_channel_credentials(root_certificates=ca_cert)
            return grpc.secure_channel(self.url, credentials)
        else:
            logger.warning(f"Creating insecure channel to {self.url}")
            # Insecure
            return grpc.insecure_channel(self.url)

    def get_metadata(self) -> list[tuple[str, str]]:
        """Get metadata (headers) for gRPC calls.

        Note: With mTLS auth, metadata is typically empty since authentication
        is done via client certificates, not headers.
        """
        return self.auth_hook.get_metadata()


# Global configuration instance
config = ObservabilityConfig()


def _get_channel() -> grpc.Channel:
    """Get or create a shared gRPC channel for observability API."""
    global _cached_channel

    if _cached_channel is not None:
        return _cached_channel

    with _channel_lock:
        # Double-check after acquiring lock
        if _cached_channel is not None:
            return _cached_channel

        _cached_channel = config.create_channel()
        return _cached_channel


def get_current_input_ids() -> list[str]:
    """
    Get the input message IDs from the current OTel span.

    This provides the context of which messages were inputs to the
    current operation, useful for backward slicing and authorization.

    Returns:
        List of input message IDs if available, empty list otherwise.
    """
    try:
        from instrumentation.otel import get_current_input_ids as _get_input_ids
        return _get_input_ids()
    except ImportError:
        return []


def get_current_message_id() -> Optional[str]:
    """
    Get the most recent input message ID from the current OTel span.

    This is the primary message in the current context, typically the
    last message that was processed before the current operation.

    Returns:
        The most recent input message ID, or None if not available.
    """
    input_ids = get_current_input_ids()
    return input_ids[-1] if input_ids else None


def configure(
    url: Optional[str] = None,
    cert_path: Optional[str] = None,
    key_path: Optional[str] = None,
    ca_path: Optional[str] = None,
    auth_hook: Optional[AuthHook] = None,
) -> None:
    """
    Configure global settings for the observability API.

    Args:
        url: gRPC server URL (e.g., "localhost:10085")
        cert_path: Path to client certificate for mTLS
        key_path: Path to client private key for mTLS
        ca_path: Path to CA certificate for server verification
        auth_hook: Auth hook for JWT/OIDC/API key authentication
    """
    global config
    if url is not None:
        config.url = url
    if cert_path is not None:
        config.cert_path = cert_path
    if key_path is not None:
        config.key_path = key_path
    if ca_path is not None:
        config.ca_path = ca_path
    if auth_hook is not None:
        config.auth_hook = auth_hook
        logger.info(f"Configured auth_hook: {type(auth_hook).__name__}")


# TODO: potentially replace the internal proto types with
# Pydantic types for user-facing functions
def record_events(events: list[Event]) -> list[str]:
    """
    Records events, corresponding to messages sent in a multi-agent system.

    Args:
        events: a list of events in the execution of the system, each corresponding
            to a message, optionally with metadata

    Returns:
        The event IDs (either provided with the events or generated by the observability
        server)
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    response = stub.RegisterEvents(
        Events(events=events),
        metadata=metadata,
    )
    return list(response.ids)


def record_dependencies(edges: list[Edge]) -> None:
    """
    Records dependencies between events.

    Args:
        dependencies: Edge data with a source and destination and optional metadata
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    stub.RegisterDependencies(
        Dependencies(edges=edges),
        metadata=metadata,
    )


def backward_slice(
    event_id: str,
    max_depth: Optional[int] = None,
) -> nwx.DiGraph:
    """
    Queries the observability server for the subgraph of the event graph containing
    the ancestors of the node with `event_id`, up to a maximum depth.

    Note: depth is computed via any direct dependency, not just proximal dependencies
    (e.g. an LLM message depends proximally on its predecessor but depends directly
    on all messages in the context window).

    Args:
        event_id: The ID of the event (as provided by `record_events()`)
        max_depth: The maximum number of steps in the dependency graph

    Returns:
        nwx.DiGraph: a directed graph of the event's ancestors containing all node
            and edge metadata.
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    graph = stub.BackwardSlice(
        SliceRequest(event_id=event_id, max_depth=max_depth),
        metadata=metadata,
    )

    return to_digraph(graph)


def forward_slice(
    event_id: str,
    max_depth: Optional[int] = None,
) -> nwx.DiGraph:
    """
    Queries the observability server for the subgraph of the event graph containing
    the descendants of the node with `event_id`, up to a maximum depth.

    Note: depth is computed via any direct dependency, not just proximal dependencies
    (e.g. an LLM message depends proximally on its predecessor but depends directly
    on all messages in the context window).

    Args:
        event_id: The ID of the event (as provided by `record_events()`)
        max_depth: The maximum number of steps in the dependency graph

    Returns:
        nwx.DiGraph: a directed graph of the event's descendants containing all node
            and edge metadata.
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    graph = stub.ForwardSlice(
        SliceRequest(event_id=event_id, max_depth=max_depth),
        metadata=metadata,
    )

    return to_digraph(graph)


def register_computations(computations: list[Computation]) -> list[str]:
    """
    Records computations (OpenTelemetry spans) in the observability graph.

    This is typically called automatically by the Neo4jSpanExporter, but can be
    used directly for manual span registration.

    Args:
        computations: A list of Computation objects representing OTel spans

    Returns:
        The span_ids of the recorded computations
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    response = stub.RegisterComputations(
        Computations(computations=computations),
        metadata=metadata,
    )
    return list(response.ids)


def get_trace(
    trace_id: str,
    start_time_ns: Optional[int] = None,
    end_time_ns: Optional[int] = None,
) -> TraceGraph:
    """
    Retrieves a full trace with all spans, messages, and their relationships.

    Args:
        trace_id: The 32-character hex trace ID
        start_time_ns: Optional start time filter (nanoseconds since epoch)
        end_time_ns: Optional end time filter (nanoseconds since epoch)

    Returns:
        TraceGraph containing:
        - computations: List of spans in the trace
        - messages: List of messages linked to the spans
        - computation_edges: CHILD_OF relationships between spans
        - message_edges: PRODUCES/CONSUMES relationships between spans and messages
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    request = TraceRequest(
        trace_id=trace_id,
        start_time_ns=start_time_ns,
        end_time_ns=end_time_ns,
    )
    return stub.GetTrace(request, metadata=metadata)


def get_span(
    span_id: str,
    include_children: bool = False,
    include_messages: bool = False,
) -> Computation:
    """
    Retrieves a single span by its span_id.

    Args:
        span_id: The 16-character hex span ID
        include_children: Whether to include child spans (not yet implemented)
        include_messages: Whether to include linked messages (not yet implemented)

    Returns:
        Computation object containing the span data
    """
    channel = _get_channel()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    request = SpanRequest(
        span_id=span_id,
        include_children=include_children,
        include_messages=include_messages,
    )
    return stub.GetSpan(request, metadata=metadata)


# Async API

_cached_async_channel: Optional[grpc.aio.Channel] = None


def _get_channel_async() -> grpc.aio.Channel:
    """Get or create a shared async gRPC channel for observability API."""
    global _cached_async_channel

    if _cached_async_channel is not None:
        return _cached_async_channel

    ca_cert = config.load_ca_cert()
    client_cert = config.load_cert()
    client_key = config.load_key()

    if client_cert and client_key and ca_cert:
        credentials = grpc.ssl_channel_credentials(
            root_certificates=ca_cert,
            private_key=client_key,
            certificate_chain=client_cert,
        )
        _cached_async_channel = grpc.aio.secure_channel(config.url, credentials)
    elif ca_cert:
        credentials = grpc.ssl_channel_credentials(root_certificates=ca_cert)
        _cached_async_channel = grpc.aio.secure_channel(config.url, credentials)
    else:
        _cached_async_channel = grpc.aio.insecure_channel(config.url)

    return _cached_async_channel


async def record_events_async(events: list[Event]) -> list[str]:
    """
    Records events asynchronously.

    Args:
        events: a list of events in the execution of the system

    Returns:
        The event IDs
    """
    channel = _get_channel_async()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    response = await stub.RegisterEvents(
        Events(events=events),
        metadata=metadata,
    )
    return list(response.ids)


async def record_dependencies_async(edges: list[Edge]) -> None:
    """
    Records dependencies between events asynchronously.

    Args:
        edges: Edge data with a source and destination and optional metadata
    """
    channel = _get_channel_async()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    await stub.RegisterDependencies(
        Dependencies(edges=edges),
        metadata=metadata,
    )


async def register_computations_async(computations: list[Computation]) -> list[str]:
    """
    Records computations (OpenTelemetry spans) asynchronously.

    Args:
        computations: A list of Computation objects representing OTel spans

    Returns:
        The span_ids of the recorded computations
    """
    channel = _get_channel_async()
    metadata = config.get_metadata()

    stub = observability_grpc.ObservabilityStub(channel)  # type: ignore
    response = await stub.RegisterComputations(
        Computations(computations=computations),
        metadata=metadata,
    )
    return list(response.ids)
