"""
Noveum Trace SDK - Cloud-first, flexible tracing for LLM applications.

This package provides observability for LLM applications and multi-agent
systems through context managers, manual spans, and integrations.

Import patterns (recommended: package root):

    >>> from noveum_trace import init, trace_context, trace_llm_call, NoveumClient
    >>> from noveum_trace import trace_operation, trace_agent_operation

    >>> import noveum_trace
    >>> noveum_trace.init(project="my-app")

Submodule imports (when needed):

    >>> from noveum_trace.integrations.langchain import NoveumTraceCallbackHandler
    >>> from noveum_trace.core.client import NoveumClient

For complete documentation, see README.md.

Example (LLM tracing with context managers):

    >>> def process_query(query: str) -> str:
    ...     cleaned_query = clean_query(query)
    ...     with noveum_trace.trace_llm_call(model="gpt-4") as span:
    ...         response = openai.chat.completions.create(...)
    ...         span.set_attribute("llm.tokens", response.usage.total_tokens)
    ...     return format_response(response)

    >>> from noveum_trace import trace_llm_call
    >>> with trace_llm_call(model="gpt-4") as span:
    ...     response = openai.chat.completions.create(...)

"""

__version__ = "1.5.21"
__author__ = "Noveum Team"
__email__ = "engineering@noveum.ai"
__license__ = "Apache-2.0"

import threading
import warnings
from typing import Any, Optional

# Agent imports
from noveum_trace.agents import (
    AgentGraph,
    AgentNode,
    AgentWorkflow,
    cleanup_agent,
    cleanup_agent_graph,
    cleanup_agent_workflow,
    cleanup_all_registries,
    cleanup_by_ttl,
    create_agent,
    create_agent_graph,
    create_agent_workflow,
    enforce_size_limits,
    get_agent,
    get_agent_graph,
    get_agent_workflow,
    get_registry_stats,
    temporary_agent_context,
)
from noveum_trace.agents import trace_agent_operation as trace_agent_op

# Context manager imports
from noveum_trace.context_managers import (
    create_child_span,
)
from noveum_trace.context_managers import trace_agent as trace_agent_operation
from noveum_trace.context_managers import (
    trace_batch_operation,
)
from noveum_trace.context_managers import trace_llm as trace_llm_call
from noveum_trace.context_managers import (
    trace_operation,
    trace_pipeline_stage,
)

# Core imports
from noveum_trace.core.client import NoveumClient
from noveum_trace.core.config import DEFAULT_ENDPOINT, configure, get_config
from noveum_trace.core.context import (
    ContextualTrace,
    get_current_span,
    get_current_trace,
    trace_context,
)
from noveum_trace.core.span import Span
from noveum_trace.core.trace import Trace

# Streaming imports
from noveum_trace.streaming import (
    create_anthropic_streaming_callback,
    create_openai_streaming_callback,
    streaming_llm,
    trace_streaming,
)

# Thread imports
from noveum_trace.threads import (
    ThreadContext,
    create_thread,
    delete_thread,
    get_thread,
    list_threads,
    trace_thread_llm,
)

# Utility imports
from noveum_trace.utils.exceptions import (
    ConfigurationError,
    InitializationError,
    InstrumentationError,
    NoveumTraceError,
    TracingError,
    TransportError,
)

# Global client instance
_client: Optional[NoveumClient] = None
_client_lock = threading.Lock()


def init(
    project: Optional[str] = None,
    api_key: Optional[str] = None,
    endpoint: Optional[str] = None,
    environment: Optional[str] = None,
    service_version: Optional[str] = None,
    organization: Optional[str] = None,
    policies: Optional[list[Any]] = None,
    guard_enabled: bool = False,
    guard_fail_open_on_backend_unavailable: bool = False,
    **kwargs: Any,
) -> None:
    """
    Initialize the Noveum Trace SDK.

    This function sets up the global tracing client.

    Args:
        project: Project name for organizing traces
        api_key: Noveum API key (defaults to NOVEUM_API_KEY env var)
        endpoint: API endpoint (defaults to https://api.noveum.ai/api)
        environment: Environment name (dev, staging, prod)
        service_version: Service/application version string (e.g. "v1.0.0").
            Exported as ``service_version`` on each trace. Defaults to the
            NOVEUM_SERVICE_VERSION env var.
        guard_fail_open_on_backend_unavailable: When the Guard control plane is
            unreachable, block every guarded call (fail closed, default) or keep
            enforcing the last-known policies (fail open, ``True``). Fail open
            trades strict enforcement for availability so a transient
            control-plane blip does not become a full LLM outage.
        **kwargs: Additional configuration options including:
                 - transport_config: Transport layer configuration
                 - tracing_config: Tracing behavior configuration
                 - security_config: Security settings configuration
                 - integrations_config: Integration settings

    Example:
        >>> import noveum_trace
        >>> noveum_trace.init(
        ...     project="my-llm-app",
        ...     environment="production",
        ...     service_version="v1.0.0",
        ...     transport_config={
        ...         "batch_size": 10,
        ...         "batch_timeout": 1.0
        ...     }
        ... )
    """
    global _client

    with _client_lock:
        # Check if client is already initialized
        if _client is not None:
            return

        # Configure the SDK - handle endpoint explicitly to ensure defaults work
        config = {}
        if project is not None:
            config["project"] = project
        if api_key is not None:
            config["api_key"] = api_key
        # Always include endpoint in config to ensure consistent behavior
        # If None, use the default endpoint
        config["endpoint"] = endpoint if endpoint is not None else DEFAULT_ENDPOINT
        if environment is not None:
            config["environment"] = environment
        if service_version is not None:
            config["service_version"] = service_version

        # Handle special config parameter mappings
        config_mappings = {
            "transport_config": "transport",
            "tracing_config": "tracing",
            "security_config": "security",
            "integrations_config": "integrations",
        }

        # Process kwargs and map config parameters to correct keys
        for key, value in kwargs.items():
            if key in config_mappings:
                # Map transport_config -> transport, etc.
                config[config_mappings[key]] = value
            else:
                # Pass through other kwargs as-is
                config[key] = value

        configure(config)

        # Initialize the client
        _client = NoveumClient()

        # Register with core module to avoid circular imports
        from noveum_trace.core import _register_client

        _register_client(_client)

        # Bootstrap guard if requested
        if guard_enabled or policies:
            import uuid

            # Warn when caller passes explicit policies alongside guard_enabled.
            # The backend poller is now the single source of truth; local policy
            # definitions may conflict with backend-fetched ones for the same name.
            if policies and guard_enabled:
                warnings.warn(
                    "Passing 'policies' together with 'guard_enabled=True' is "
                    "deprecated. When guard_enabled=True the backend poller "
                    "becomes the single source of truth for policy definitions. "
                    "Remove the 'policies' argument and configure policies via "
                    "the Noveum dashboard instead.",
                    DeprecationWarning,
                    stacklevel=2,
                )

            try:
                from noveum_trace.guard._state import set_guard
                from noveum_trace.guard.engine import PolicyEngine
                from noveum_trace.guard.poller import PolicyPoller
                from noveum_trace.guard.types import PolicyContext

                # A real API key selects the HTTP backend (spend is shared across
                # processes via /policies/state + /policies/usage). No key keeps
                # the in-memory stub for local dev / tests. Read from the
                # resolved SDK config so NOVEUM_API_KEY (not just the explicit
                # api_key argument) is honored.
                _resolved_key = get_config().api_key or ""
                if _resolved_key:
                    from noveum_trace.guard.api_client_http import HttpGuardAPIClient

                    _api_client: Any = HttpGuardAPIClient(
                        api_key=_resolved_key,
                        base_url=endpoint or DEFAULT_ENDPOINT,
                        organization_slug=organization,
                    )
                else:
                    from noveum_trace.guard.api_client import GuardAPIClient

                    _api_client = GuardAPIClient(
                        api_key="",
                        base_url=endpoint or DEFAULT_ENDPOINT,
                    )
                _engine = PolicyEngine(
                    _api_client,
                    fail_open_on_backend_unavailable=(
                        guard_fail_open_on_backend_unavailable
                    ),
                )
                # Use a unique sentinel when project is not provided so that
                # different uninamed callers do not share the same spend scope.
                _project_id = (
                    project if project is not None else f"__unscoped_{uuid.uuid4()}"
                )
                _ctx = PolicyContext(
                    project_id=_project_id,
                    organization_id=None,
                    environment=environment or "production",
                    trace_id=None,
                    span_id=None,
                    call_id=str(uuid.uuid4()),
                )
                # Bind the ambient context before attaching so each policy can adopt
                # the project scope its background poll() needs.
                for _policy in policies or []:
                    _policy.bind_context(_ctx)
                    _engine.attach(_policy)
                # Pass project_id and context explicitly: start() runs before
                # set_guard() below, so the poller cannot read either from _state.
                _poller = PolicyPoller(_engine, project_id=_project_id, context=_ctx)
                try:
                    _poller.start()
                except Exception:
                    # Don't leave the poller half-started; guard stays uninitialized.
                    _poller.stop()
                    raise
                set_guard(_engine, _ctx, _poller)
            except Exception:
                # Roll back the global client registration so the SDK is not
                # left in a partially-initialized state. A subsequent init()
                # call will retry the full bootstrap.
                _client = None
                from noveum_trace.core import _unregister_client

                _unregister_client()
                raise


def shutdown() -> None:
    """
    Shutdown the Noveum Trace SDK.

    This function flushes any pending traces and cleanly shuts down
    the tracing client.
    """
    global _client
    with _client_lock:
        # Stop the guard poller thread and clear guard state if running
        from noveum_trace.guard import _state as _guard_state

        _guard_poller = _guard_state.get_poller()
        if _guard_poller is not None:
            _guard_poller.stop()
        # Drain any queued usage pushes before dropping the engine.
        _guard_engine = _guard_state.get_engine()
        if _guard_engine is not None:
            _close = getattr(_guard_engine._api_client, "close", None)
            if callable(_close):
                _close()
        _guard_state.clear()

        if _client:
            _client.shutdown()
            _client = None

            # Unregister from core module
            from noveum_trace.core import _unregister_client

            _unregister_client()


def flush() -> None:
    """
    Flush any pending traces to the Noveum platform.

    This function blocks until all pending traces have been sent
    or the flush timeout is reached.
    """
    global _client
    with _client_lock:
        if _client:
            _client.flush()


def get_client() -> NoveumClient:
    """
    Get the global Noveum client instance.

    Returns:
        The global NoveumClient instance

    Raises:
        InitializationError: If the SDK has not been initialized
    """
    global _client
    with _client_lock:
        if _client is None:
            raise InitializationError(
                "Noveum Trace SDK not initialized. Call noveum_trace.init() first."
            )
        return _client


def is_initialized() -> bool:
    """
    Check if the Noveum Trace SDK has been initialized.

    Returns:
        True if initialized, False otherwise
    """
    global _client
    with _client_lock:
        return _client is not None


# Plugin system
def register_plugin(plugin: Any) -> None:
    """
    Register a custom plugin with the Noveum Trace SDK.

    Args:
        plugin: Plugin instance implementing the BasePlugin interface
    """
    # from noveum_trace.plugins import register_plugin as _register_plugin
    # _register_plugin(plugin)
    raise NotImplementedError("Plugin system not yet implemented")


def list_plugins() -> list[str]:
    """
    List all registered plugins.

    Returns:
        List of registered plugin names
    """
    # from noveum_trace.plugins import list_plugins as _list_plugins
    # return _list_plugins()
    raise NotImplementedError("Plugin system not yet implemented")


# Convenience functions for manual instrumentation
def start_trace(name: str, **kwargs: Any) -> "ContextualTrace":
    """
    Manually start a new trace.

    Args:
        name: Trace name
        **kwargs: Additional trace attributes

    Returns:
        New ContextualTrace instance that can be used as a context manager
    """
    client = get_client()
    return client.create_contextual_trace(name, **kwargs)


def start_span(name: str, **kwargs: Any) -> Span:
    """
    Manually start a new span.

    Args:
        name: Span name
        **kwargs: Additional span attributes

    Returns:
        New Span instance
    """
    client = get_client()
    return client.start_span(name, **kwargs)


# Integrations imports (conditional)
_integration_exports = []
try:
    from noveum_trace.integrations.langchain import NoveumTraceCallbackHandler

    _integration_exports.append("NoveumTraceCallbackHandler")
except ImportError:
    # LangChain not installed
    pass

try:
    from noveum_trace.integrations.livekit import (
        LiveKitSTTWrapper,
        LiveKitTTSWrapper,
    )

    _integration_exports.extend(["LiveKitSTTWrapper", "LiveKitTTSWrapper"])
except ImportError:
    # LiveKit not installed
    pass

try:
    from noveum_trace.integrations.crewai import (
        NoveumCrewAIListener,
        setup_crewai_tracing,
    )

    _integration_exports.extend(["NoveumCrewAIListener", "setup_crewai_tracing"])
except ImportError:
    # CrewAI not installed
    pass

# Export public API
__all__ = [
    # Core functions
    "init",
    "shutdown",
    "flush",
    "configure",
    "get_config",
    "get_client",
    "is_initialized",
    # Context managers
    "trace_context",
    "trace_llm_call",
    "trace_agent_operation",
    "trace_operation",
    "trace_batch_operation",
    "trace_pipeline_stage",
    "create_child_span",
    # Streaming
    "trace_streaming",
    "streaming_llm",
    "create_openai_streaming_callback",
    "create_anthropic_streaming_callback",
    # Threads
    "create_thread",
    "get_thread",
    "delete_thread",
    "list_threads",
    "trace_thread_llm",
    "ThreadContext",
    # Agents
    "create_agent",
    "get_agent",
    "create_agent_graph",
    "get_agent_graph",
    "create_agent_workflow",
    "get_agent_workflow",
    "trace_agent_op",
    "AgentNode",
    "AgentGraph",
    "AgentWorkflow",
    # Agent cleanup
    "cleanup_agent",
    "cleanup_agent_graph",
    "cleanup_agent_workflow",
    "cleanup_all_registries",
    "cleanup_by_ttl",
    "enforce_size_limits",
    "get_registry_stats",
    "temporary_agent_context",
    # Context management
    "get_current_trace",
    "get_current_span",
    # Manual instrumentation
    "start_trace",
    "start_span",
    # Core classes
    "Trace",
    "Span",
    "NoveumClient",
    # Plugin system
    "register_plugin",
    "list_plugins",
    # Integrations (conditional)
    *_integration_exports,
    # Exceptions
    "NoveumTraceError",
    "ConfigurationError",
    "TransportError",
    "InstrumentationError",
    "TracingError",
    # Metadata
    "__version__",
    "__author__",
    "__email__",
    "__license__",
]