"""
Agnost Conversation SDK

A Python SDK for tracking AI conversations and events.

Example:
    import agnost_conversation as agnost

    # Initialize with your write key
    agnost.init("your-write-key")

    # Track an AI interaction
    agnost.track_ai(
        user_id="user_123",
        event="chat_completion",
        input="What is the weather?",
        output="The weather is sunny.",
        agent_name="weather-agent"
    )

    # Identify a user
    agnost.identify("user_123", {
        "name": "John Doe",
        "email": "john@example.com",
        "plan": "paid"
    })

    # For partial events
    interaction = agnost.begin(
        user_id="user_123",
        event="chat_completion"
    )
    interaction.set_input("Hello")
    interaction.finish(output="Hi there!")

    # Ensure events are sent before exit
    agnost.shutdown()
"""

from .client import _client, ConversationClient
from .types import (
    ConversationConfig,
    Interaction,
    EventProperties,
    UserTraits,
    EventData,
    EventIOData,
    ConversationData,
)

__version__ = "0.1.0"
__all__ = [
    # Module-level functions (use the global client)
    "init",
    "track_ai",
    "identify",
    "begin",
    "flush",
    "shutdown",
    "set_debug_logs",
    # Classes for direct instantiation
    "ConversationClient",
    "ConversationConfig",
    "Interaction",
    # Type aliases
    "EventProperties",
    "UserTraits",
    "EventData",
    "EventIOData",
    "ConversationData",
]


def init(write_key: str, config: ConversationConfig = None, **kwargs) -> bool:
    """
    Initialize the Agnost Conversation SDK.

    Args:
        write_key: Your Agnost organization/write key
        config: Optional ConversationConfig instance
        **kwargs: Additional config options:
            - endpoint: API endpoint (default: "https://api.agnost.ai")
            - debug: Enable debug logging (default: False)

    Returns:
        bool: True if initialization was successful

    Example:
        import agnost_conversation as agnost

        # Basic initialization
        agnost.init("your-write-key")

        # With custom config
        agnost.init("your-write-key", endpoint="https://custom.endpoint.com", debug=True)
    """
    return _client.init(write_key, config, **kwargs)


def track_ai(
    user_id: str,
    event: str,
    input: str = None,
    output: str = None,
    agent_name: str = None,
    convo_id: str = None,
    properties: dict = None,
    timestamp=None,
    success: bool = True,
    latency: int = None,
) -> str:
    """
    Track an AI interaction event.

    Args:
        user_id: User identifier (required)
        event: Event name (required)
        input: Input text/prompt
        output: Output/response text
        agent_name: Agent name (e.g., "weather-agent", "support-bot")
        convo_id: Conversation ID for grouping related events
        properties: Additional event properties
        timestamp: Event timestamp (defaults to now)
        success: Whether the interaction was successful
        latency: Execution time in milliseconds

    Returns:
        str: Event ID if queued successfully, None otherwise

    Example:
        event_id = agnost.track_ai(
            user_id="user_123",
            event="chat_completion",
            input="What is the weather?",
            output="The weather is sunny.",
            agent_name="weather-agent",
            convo_id="conv_abc",
            properties={"temperature": "0.7"},
            latency=1500
        )
    """
    return _client.track_ai(
        user_id=user_id,
        event=event,
        input=input,
        output=output,
        agent_name=agent_name,
        convo_id=convo_id,
        properties=properties,
        timestamp=timestamp,
        success=success,
        latency=latency,
    )


def identify(user_id: str, traits: dict) -> bool:
    """
    Associate user characteristics with a user ID.

    Args:
        user_id: User identifier
        traits: User traits dictionary

    Returns:
        bool: True if identification was successful

    Example:
        agnost.identify("user_123", {
            "name": "John Doe",
            "email": "john@example.com",
            "age": 30,
            "plan": "paid"
        })
    """
    return _client.identify(user_id, traits)


def begin(
    user_id: str,
    event: str,
    agent_name: str = None,
    convo_id: str = None,
    properties: dict = None,
) -> Interaction:
    """
    Begin a partial interaction for deferred completion.

    Use this when you want to track input/output separately or when
    the output is not immediately available.

    Args:
        user_id: User identifier
        event: Event name
        agent_name: Agent name
        convo_id: Conversation ID
        properties: Initial properties

    Returns:
        Interaction: An interaction object to complete later

    Example:
        interaction = agnost.begin(
            user_id="user_123",
            event="chat_completion",
            agent_name="weather-agent"
        )
        interaction.set_input("What is the weather?")
        # ... do async work ...
        interaction.finish(output="The weather is sunny.", latency=1500)
    """
    return _client.begin(
        user_id=user_id,
        event=event,
        agent_name=agent_name,
        convo_id=convo_id,
        properties=properties,
    )


def flush() -> None:
    """
    Manually flush all queued events.

    Call this before your application exits if you want to ensure
    all events are sent.

    Example:
        agnost.flush()
    """
    _client.flush()


def shutdown() -> None:
    """
    Shutdown the SDK and flush remaining events.

    Call this when your application is shutting down to ensure
    all events are sent and resources are cleaned up.

    Example:
        agnost.shutdown()
    """
    _client.shutdown()


def set_debug_logs(enabled: bool) -> None:
    """
    Enable or disable debug logging.

    Args:
        enabled: True to enable debug logs, False to disable

    Example:
        agnost.set_debug_logs(True)  # See queued events
    """
    _client.set_debug_logs(enabled)
