from __future__ import annotations

from typing import Any, Optional
from .transport import configure, post
from .trace import Trace


class TokenCat:
    _initialized: bool = False

    @classmethod
    def init(cls, *, api_key: str, base_url: Optional[str] = None) -> None:
        configure(api_key, base_url)
        cls._initialized = True

    @classmethod
    async def start_trace(
        cls,
        name: str,
        *,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        input: Any = None,
        metadata: Optional[dict[str, Any]] = None,
        prompt_id: Optional[str] = None,
        prompt_version: Optional[int] = None,
    ) -> Trace:
        if not cls._initialized:
            print("[ZappyBee] SDK not initialized. Call ZappyBee.init() (or TokenCat.init()) first.")
            return Trace("unknown")

        meta = dict(metadata or {})
        if prompt_id:
            meta["prompt_id"] = prompt_id
        if prompt_version:
            meta["prompt_version"] = prompt_version

        # Important: omit optional keys when they're None.
        # The backend expects absent keys for optionals, not explicit nulls.
        body: dict[str, Any] = {"name": name, "input": input, "metadata": meta}
        if user_id is not None:
            body["user_id"] = user_id
        if session_id is not None:
            body["session_id"] = session_id

        result = await post("/api/v1/traces", body)

        trace_id = result.get("trace", {}).get("id", "unknown") if result else "unknown"
        return Trace(trace_id)

    @classmethod
    def wrap(cls, client: Any) -> Any:
        """Auto-detect and wrap a supported client (Anthropic or OpenAI)."""
        client_module = type(client).__module__ or ""
        if "anthropic" in client_module:
            from .wrappers.anthropic import wrap_anthropic
            return wrap_anthropic(client)
        if "openai" in client_module:
            from .wrappers.openai import wrap_openai
            return wrap_openai(client)
        # Fallback: try Anthropic pattern first, then OpenAI
        if hasattr(client, "messages"):
            from .wrappers.anthropic import wrap_anthropic
            return wrap_anthropic(client)
        if hasattr(client, "chat"):
            from .wrappers.openai import wrap_openai
            return wrap_openai(client)
        print("[ZappyBee] Unknown client type. Use wrap_anthropic(), wrap_openai(), or wrap_crewai() directly.")
        return client
