"""Bandito SDK - trace LLM applications."""

from __future__ import annotations

from typing import Any

from .context import TraceContext, flush
from .models import Evaluation, Span, Trace
from .storage import get_store

__all__ = [
    "trace",
    "flush",
    "Trace",
    "Span",
    "Evaluation",
    "TraceContext",
    "get_store",
]

from .storage.base import TraceReader as _TraceReader

_store_cache: dict[tuple, _TraceReader] = {}


def _get_store(project: str | None = None, force_local: bool = False) -> _TraceReader:
    key = (project, force_local)
    if key not in _store_cache:
        _store_cache[key] = get_store(
            project=project,
            backend="local" if force_local else None,
        )
    return _store_cache[key]


def trace(
    project: str,
    *,
    input: Any = None,
    tag: str | None = None,
    session_id: str | None = None,
    user_id: str | None = None,
    metadata: dict[str, Any] | None = None,
    tags: list[str] | None = None,
    local: bool = False,
    storage: bool = True,
    writer: Any = None,
) -> TraceContext:
    import os

    resolved_tag = tag or os.environ.get("BANDITO_TAG")
    if resolved_tag:
        tags = list(tags or [])
        if resolved_tag not in tags:
            tags.insert(0, resolved_tag)

    if not storage:
        writer = None
    elif writer is None:
        if local or os.environ.get("BANDITO_LOCAL", "").lower() in ("true", "1"):
            writer = _get_store(project, force_local=True)
        else:
            writer = _get_store(project)

    kwargs: dict[str, Any] = {}
    if input is not None:
        kwargs["input"] = input
    if session_id is not None:
        kwargs["session_id"] = session_id
    if user_id is not None:
        kwargs["user_id"] = user_id
    if metadata is not None:
        kwargs["metadata"] = metadata
    if tags is not None:
        kwargs["tags"] = tags
    return TraceContext(project, writer=writer, **kwargs)
