"""Staso Python SDK – observe, enforce, evaluate, and debug AI agents."""
from __future__ import annotations

from collections.abc import Generator
from contextlib import contextmanager

from . import integrations
from ._version import __version__
from .client import Client, _set_client, get_client
from .config import Config
from .decorators import agent, tool
from .decorators import trace_fn as trace
from .provider import session_id_var

__all__ = [
    "__version__",
    "init",
    "shutdown",
    "agent",
    "tool",
    "trace",
    "session",
    "integrations",
    "Client",
    "Config",
    "get_client",
]


def init(
    *,
    api_key: str,
    agent_id: str,
    base_url: str = "http://localhost:6999",
    batch_size: int = 100,
    flush_interval: float = 5.0,
    max_queue_size: int = 10_000,
    enabled: bool = True,
    debug: bool = False,
    environment: str = "default",
    workspace_slug: str = "",
) -> Client:
    """Initialise the global Staso client.

    Must be called before any decorator or integration is used.

    >>> import staso as st
    >>> al.init(api_key="ak_...", agent_id="my-agent")
    """
    cfg = Config(
        api_key=api_key,
        agent_id=agent_id,
        base_url=base_url,
        batch_size=batch_size,
        flush_interval=flush_interval,
        max_queue_size=max_queue_size,
        enabled=enabled,
        debug=debug,
        environment=environment,
        workspace_slug=workspace_slug,
    )
    client = Client(cfg)
    _set_client(client)
    return client


def shutdown() -> None:
    """Flush pending spans and tear down the background transport."""
    client = get_client()
    if client:
        client.shutdown()


@contextmanager
def session(session_id: str) -> Generator[None, None, None]:
    """Context manager that sets the *session_id* for all spans in its scope.

    >>> with al.session("conv-42"):
    ...     run_agent(query)
    """
    token = session_id_var.set(session_id)
    try:
        yield
    finally:
        session_id_var.reset(token)
