"""
Amplitude AI SDK - LLM usage tracking for Amplitude Analytics

Provides automatic tracking for OpenAI, Anthropic, Google Gemini, and other LLM providers.

Quick start::

    from amplitude import Amplitude
    from amplitude_ai import AmplitudeAI

    amplitude = Amplitude("api-key")
    ai = AmplitudeAI(amplitude=amplitude)
"""

import importlib
from typing import Any

from .client import AmplitudeAI, BoundAgent, TenantHandle, Session
from .config import AIConfig, ContentMode
from .core.decorators import ToolCallTracker, tool, observe
from .patching import (
    patch,
    patch_anthropic,
    patch_azure_openai,
    patch_gemini,
    patch_mistral,
    patch_openai,
    unpatch,
    unpatch_anthropic,
    unpatch_azure_openai,
    unpatch_gemini,
    unpatch_mistral,
    unpatch_openai,
)
from .wrappers import AmplitudeAIWrapError, wrap
from .core.enrichments import EvidenceQuote, MessageLabel, RubricScore, SessionEnrichments, TopicClassification
from .exceptions import AmplitudeAIError, ConfigurationError, ProviderError, TrackingError, ValidationError
from .testing import MockAmplitudeAI
from .core.tracking import (
    # Event type constants
    EVENT_AI_RESPONSE,
    EVENT_EMBEDDING,
    EVENT_SCORE,
    EVENT_SESSION_END,
    EVENT_SESSION_ENRICHMENT,
    EVENT_SPAN,
    EVENT_TOOL_CALL,
    EVENT_USER_MESSAGE,
    # Property name constants
    PROP_AGENT_ID,
    PROP_AGENT_VERSION,
    PROP_CACHE_CREATION_TOKENS,
    PROP_CACHE_READ_TOKENS,
    PROP_COMMENT,
    PROP_COMPONENT_TYPE,
    PROP_CONTEXT,
    PROP_COST_USD,
    PROP_CUSTOMER_ORG_ID,
    PROP_EMBEDDING_DIMENSIONS,
    PROP_ENRICHMENTS,
    PROP_ENV,
    PROP_ERROR_MESSAGE,
    PROP_FINISH_REASON,
    PROP_HAS_REASONING,
    PROP_INPUT_STATE,
    PROP_INPUT_TOKENS,
    PROP_INVOCATION_ID,
    PROP_IS_ERROR,
    PROP_LATENCY_MS,
    PROP_LOCALE,
    PROP_MESSAGE_ID,
    PROP_MODEL_NAME,
    PROP_OUTPUT_STATE,
    PROP_OUTPUT_TOKENS,
    PROP_PARENT_AGENT_ID,
    PROP_PARENT_MESSAGE_ID,
    PROP_PARENT_SPAN_ID,
    PROP_PROVIDER,
    PROP_REASONING_CONTENT,
    PROP_REASONING_TOKENS,
    PROP_RUNTIME,
    PROP_SCORE_NAME,
    PROP_EVALUATION_SOURCE,
    PROP_SCORE_VALUE,
    PROP_SDK_VERSION,
    PROP_SESSION_ID,
    PROP_SPAN_ID,
    PROP_SPAN_KIND,
    PROP_SPAN_NAME,
    PROP_TARGET_ID,
    PROP_TARGET_TYPE,
    PROP_TOOL_CALLS,
    PROP_TOOL_INPUT,
    PROP_TOOL_NAME,
    PROP_TOOL_OUTPUT,
    PROP_TOOL_SUCCESS,
    PROP_TOTAL_TOKENS,
    PROP_TRACE_ID,
    PROP_TTFB_MS,
    PROP_TURN_ID,
    PROP_SYSTEM_PROMPT,
    PROP_SYSTEM_PROMPT_LENGTH,
    PROP_TEMPERATURE,
    PROP_MAX_OUTPUT_TOKENS,
    PROP_TOP_P,
    PROP_IS_STREAMING,
    PROP_PROMPT_ID,
    PROP_IS_REGENERATION,
    PROP_IS_EDIT,
    PROP_EDITED_MESSAGE_ID,
    PROP_WAS_COPIED,
    PROP_ABANDONMENT_TURN,
    PROP_WAS_CACHED,
    PROP_MODEL_TIER,
    PROP_HAS_ATTACHMENTS,
    PROP_ATTACHMENT_TYPES,
    PROP_ATTACHMENT_COUNT,
    PROP_TOTAL_ATTACHMENT_SIZE,
    PROP_ATTACHMENTS,
    PROP_MESSAGE_LABELS,
    PROP_MESSAGE_LABEL_MAP,
    PROP_QUALITY_SCORE,
    PROP_SENTIMENT_SCORE,
    PROP_TASK_FAILURE_TYPE,
    PROP_TASK_FAILURE_REASON,
    PROP_AGENT_CHAIN,
    PROP_ROOT_AGENT_NAME,
    PROP_REQUEST_COMPLEXITY,
    SDK_VERSION,
    track_ai_message,
    track_conversation,
    track_embedding,
    track_score,
    track_session_end,
    track_session_enrichment,
    track_span,
    track_tool_call,
    track_user_message,
)


def _import_providers() -> tuple[Any, ...]:
    providers: dict[str, Any] = {}

    for name, module, cls_name in [
        ("OpenAI", ".providers.openai", "OpenAI"),
        ("Anthropic", ".providers.anthropic", "Anthropic"),
        ("Gemini", ".providers.gemini", "Gemini"),
        ("AzureOpenAI", ".providers.azure_openai", "AzureOpenAI"),
        ("Bedrock", ".providers.bedrock", "Bedrock"),
        ("Mistral", ".providers.mistral", "Mistral"),
    ]:
        try:
            mod = importlib.import_module(module, package=__name__)
            providers[name] = getattr(mod, cls_name)
        except (ImportError, AttributeError):
            providers[name] = None

    return (
        providers["OpenAI"],
        providers["Anthropic"],
        providers["Gemini"],
        providers["AzureOpenAI"],
        providers["Bedrock"],
        providers["Mistral"],
    )


OpenAI: Any
Anthropic: Any
Gemini: Any
AzureOpenAI: Any
Bedrock: Any
Mistral: Any
OpenAI, Anthropic, Gemini, AzureOpenAI, Bedrock, Mistral = _import_providers()


def _import_integrations() -> tuple[Any, ...]:
    results: dict[str, Any] = {}

    for name, module, cls_name in [
        ("AmplitudeCallbackHandler", ".integrations.langchain", "AmplitudeCallbackHandler"),
        ("create_amplitude_callback", ".integrations.langchain", "create_amplitude_callback"),
        ("AmplitudeSpanExporter", ".integrations.opentelemetry", "AmplitudeSpanExporter"),
        ("AmplitudeAgentExporter", ".integrations.opentelemetry", "AmplitudeAgentExporter"),
        ("AmplitudeGenAIExporter", ".integrations.opentelemetry", "AmplitudeGenAIExporter"),
        ("AmplitudeLlamaIndexHandler", ".integrations.llamaindex", "AmplitudeLlamaIndexHandler"),
        ("create_amplitude_llamaindex_handler", ".integrations.llamaindex", "create_amplitude_llamaindex_handler"),
    ]:
        try:
            mod = importlib.import_module(module, package=__name__)
            results[name] = getattr(mod, cls_name)
        except (ImportError, AttributeError):
            results[name] = None

    return (
        results["AmplitudeCallbackHandler"],
        results["create_amplitude_callback"],
        results["AmplitudeSpanExporter"],
        results["AmplitudeAgentExporter"],
        results["AmplitudeGenAIExporter"],
        results["AmplitudeLlamaIndexHandler"],
        results["create_amplitude_llamaindex_handler"],
    )


AmplitudeCallbackHandler: Any
create_amplitude_callback: Any
AmplitudeSpanExporter: Any
AmplitudeAgentExporter: Any
AmplitudeGenAIExporter: Any
AmplitudeLlamaIndexHandler: Any
create_amplitude_llamaindex_handler: Any
(
    AmplitudeCallbackHandler,
    create_amplitude_callback,
    AmplitudeSpanExporter,
    AmplitudeAgentExporter,
    AmplitudeGenAIExporter,
    AmplitudeLlamaIndexHandler,
    create_amplitude_llamaindex_handler,
) = _import_integrations()

__version__ = SDK_VERSION

__all__ = [
    # Public API (v0.3.0)
    "AIConfig",
    "AmplitudeAI",
    "AmplitudeAIError",
    "BoundAgent",
    "TenantHandle",
    "ConfigurationError",
    "ContentMode",
    "EvidenceQuote",
    "MessageLabel",
    "MockAmplitudeAI",
    "ProviderError",
    "RubricScore",
    "Session",
    "SessionEnrichments",
    "TopicClassification",
    "TrackingError",
    "ValidationError",
    # Event type constants
    "EVENT_AI_RESPONSE",
    "EVENT_EMBEDDING",
    "EVENT_SCORE",
    "EVENT_SESSION_END",
    "EVENT_SESSION_ENRICHMENT",
    "EVENT_SPAN",
    "EVENT_TOOL_CALL",
    "EVENT_USER_MESSAGE",
    # Property name constants
    "PROP_AGENT_ID",
    "PROP_AGENT_VERSION",
    "PROP_CACHE_CREATION_TOKENS",
    "PROP_CACHE_READ_TOKENS",
    "PROP_COMMENT",
    "PROP_COMPONENT_TYPE",
    "PROP_CONTEXT",
    "PROP_COST_USD",
    "PROP_CUSTOMER_ORG_ID",
    "PROP_EMBEDDING_DIMENSIONS",
    "PROP_ENRICHMENTS",
    "PROP_ENV",
    "PROP_ERROR_MESSAGE",
    "PROP_FINISH_REASON",
    "PROP_HAS_REASONING",
    "PROP_INPUT_STATE",
    "PROP_INPUT_TOKENS",
    "PROP_INVOCATION_ID",
    "PROP_IS_ERROR",
    "PROP_LATENCY_MS",
    "PROP_LOCALE",
    "PROP_MESSAGE_ID",
    "PROP_MODEL_NAME",
    "PROP_OUTPUT_STATE",
    "PROP_OUTPUT_TOKENS",
    "PROP_PARENT_AGENT_ID",
    "PROP_PARENT_MESSAGE_ID",
    "PROP_PARENT_SPAN_ID",
    "PROP_PROVIDER",
    "PROP_REASONING_CONTENT",
    "PROP_REASONING_TOKENS",
    "PROP_RUNTIME",
    "PROP_SCORE_NAME",
    "PROP_EVALUATION_SOURCE",
    "PROP_SCORE_VALUE",
    "PROP_SDK_VERSION",
    "PROP_SESSION_ID",
    "PROP_SPAN_ID",
    "PROP_SPAN_KIND",
    "PROP_SPAN_NAME",
    "PROP_TARGET_ID",
    "PROP_TARGET_TYPE",
    "PROP_TOOL_CALLS",
    "PROP_TOOL_INPUT",
    "PROP_TOOL_NAME",
    "PROP_TOOL_OUTPUT",
    "PROP_TOOL_SUCCESS",
    "PROP_TOTAL_TOKENS",
    "PROP_TRACE_ID",
    "PROP_TTFB_MS",
    "PROP_TURN_ID",
    # Model configuration constants
    "PROP_SYSTEM_PROMPT",
    "PROP_SYSTEM_PROMPT_LENGTH",
    "PROP_TEMPERATURE",
    "PROP_MAX_OUTPUT_TOKENS",
    "PROP_TOP_P",
    "PROP_IS_STREAMING",
    "PROP_PROMPT_ID",
    # Implicit feedback and attachment constants
    "PROP_IS_REGENERATION",
    "PROP_IS_EDIT",
    "PROP_EDITED_MESSAGE_ID",
    "PROP_WAS_COPIED",
    "PROP_ABANDONMENT_TURN",
    "PROP_WAS_CACHED",
    "PROP_MODEL_TIER",
    "PROP_HAS_ATTACHMENTS",
    "PROP_ATTACHMENT_TYPES",
    "PROP_ATTACHMENT_COUNT",
    "PROP_TOTAL_ATTACHMENT_SIZE",
    "PROP_ATTACHMENTS",
    "PROP_MESSAGE_LABELS",
    "PROP_MESSAGE_LABEL_MAP",
    "PROP_QUALITY_SCORE",
    "PROP_SENTIMENT_SCORE",
    "PROP_TASK_FAILURE_TYPE",
    "PROP_TASK_FAILURE_REASON",
    "PROP_AGENT_CHAIN",
    "PROP_ROOT_AGENT_NAME",
    "PROP_REQUEST_COMPLEXITY",
    "SDK_VERSION",
    # Integrations
    "AmplitudeAgentExporter",
    "AmplitudeCallbackHandler",
    "AmplitudeGenAIExporter",
    "AmplitudeLlamaIndexHandler",
    "AmplitudeSpanExporter",
    "create_amplitude_callback",
    "create_amplitude_llamaindex_handler",
    # Provider wrappers
    "Anthropic",
    "AzureOpenAI",
    "Bedrock",
    "Gemini",
    "Mistral",
    "OpenAI",
    "ToolCallTracker",
    # Tool decorator
    "tool",
    "observe",
    # Patching (zero-code instrumentation)
    "patch",
    "patch_openai",
    "patch_anthropic",
    "patch_azure_openai",
    "patch_gemini",
    "patch_mistral",
    "unpatch",
    "unpatch_openai",
    "unpatch_anthropic",
    "unpatch_azure_openai",
    "unpatch_gemini",
    "unpatch_mistral",
    # Wrap convenience
    "wrap",
    "AmplitudeAIWrapError",
    # Tracking functions (lower-level API)
    "track_ai_message",
    "track_conversation",
    "track_embedding",
    "track_score",
    "track_session_end",
    "track_session_enrichment",
    "track_span",
    "track_tool_call",
    "track_user_message",
]
