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

import os
import uuid
from collections.abc import Generator
from contextlib import contextmanager

from . import integrations
from ._version import __version__
from .client import Client, _set_client, get_client, get_tracer
from .config import Config
from .context import session_id_var, user_id_var
from .decorators import agent, tool
from .decorators import trace_fn as trace
from .provider import Span
from .annotation import AnnotationDataType, annotate
from .types import SpanKind, SpanStatus

__all__ = [
    "__version__",
    "init",
    "shutdown",
    "agent",
    "tool",
    "trace",
    "span",
    "conversation",
    "integrations",
    "Client",
    "Config",
    "Span",
    "SpanKind",
    "SpanStatus",
    "get_client",
    "annotate",
    "AnnotationDataType",
]


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

    Arguments override environment variables. Env vars checked:
    STASO_API_KEY, STASO_AGENT_ID, STASO_BASE_URL, STASO_ENVIRONMENT,
    STASO_WORKSPACE_SLUG, STASO_DEBUG
    """
    resolved_api_key = api_key or os.environ.get("STASO_API_KEY")
    resolved_agent_id = agent_id or os.environ.get("STASO_AGENT_ID")

    if not resolved_api_key:
        raise ValueError(
            "api_key is required. Pass it directly or set STASO_API_KEY env var."
        )
    if not resolved_agent_id:
        raise ValueError(
            "agent_id is required. Pass it directly or set STASO_AGENT_ID env var."
        )

    cfg = Config(
        api_key=resolved_api_key,
        agent_id=resolved_agent_id,
        base_url=base_url or os.environ.get("STASO_BASE_URL", "http://localhost:6999"),
        batch_size=batch_size,
        flush_interval=flush_interval,
        max_queue_size=max_queue_size,
        enabled=enabled,
        debug=debug if debug is not None else os.environ.get("STASO_DEBUG", "").lower() in ("1", "true"),
        environment=environment or os.environ.get("STASO_ENVIRONMENT", "default"),
        workspace_slug=workspace_slug or os.environ.get("STASO_WORKSPACE_SLUG", ""),
        capture_messages=capture_messages,
    )
    client = Client(cfg)
    _set_client(client)
    return client


@contextmanager
def span(name: str, *, kind: str | SpanKind = SpanKind.CHAIN) -> Generator[Span, None, None]:
    """Create a manual span for custom instrumentation.

    >>> with st.span("my-operation", kind="tool") as s:
    ...     s.metadata["custom"] = "value"
    ...     result = do_work()
    """
    _kind = SpanKind(kind.lower()) if isinstance(kind, str) else kind
    tracer = get_tracer()
    with tracer.start_as_current_span(name, kind=_kind) as s:
        yield s


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


@contextmanager
def conversation(
    conversation_id: str | None = None,
    *,
    user_id: str | None = None,
) -> Generator[None, None, None]:
    """Context manager that groups all spans in its scope under a conversation.

    If *conversation_id* is not provided, a new UUID is generated automatically.
    If *user_id* is provided, all spans created within this context will carry it.

    >>> with st.conversation("conv-42", user_id="user-alice"):
    ...     run_agent(query)

    >>> with st.conversation():  # auto-generated ID
    ...     run_agent(query)
    """
    sid = conversation_id or str(uuid.uuid4())
    session_token = session_id_var.set(sid)
    user_token = user_id_var.set(user_id) if user_id else None
    try:
        yield
    finally:
        session_id_var.reset(session_token)
        if user_token is not None:
            user_id_var.reset(user_token)
