One drop-in client.
Every Cartie header — automatic.

"""
Cartie SDK — Python
====================
Drop-in wrapper around the OpenAI / Anthropic Python SDKs that auto-injects:

  * X-Cartie-Tenant            (per-department attribution)
  * X-Cartie-Source-Cloud      (egress preflight)
  * X-Cartie-Dest-Cloud        (egress preflight)
  * X-Cartie-Data-GB           (egress preflight)
  * X-Cartie-Token-Budget-Pct  (egress preflight)
  * X-Cartie-Egress-Enforce    (block on rose verdict)

Usage — minimum viable (auto-attribution):

    import cartie
    cartie.set_proxy(key="ck_live_...")
    cartie.set_tenant("marketing")

    client = cartie.openai()  # exactly like openai.OpenAI()
    client.chat.completions.create(
        model="gpt-5.2",
        messages=[{"role": "user", "content": "Draft a blog headline"}],
    )

Usage — per-request egress gate (CI-blocking pipeline):

    with cartie.egress(source="gcp", dest="aws", data_gb=1024,
                       token_pct=0.05, enforce=True):
        # If verdict is rose ("MOVE THE MODEL"), proxy returns 412 BEFORE
        # the upstream call — your spend is saved.
        cartie.openai().chat.completions.create(...)

Context-manager scoping works for both tenant() and egress():

    with cartie.tenant("sales-emea"):
        ...   # all calls inside attribute to sales-emea

Author: Cartie AI
License: MIT
"""

from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Iterator, Optional

__version__ = "1.0.0"
__all__ = [
    "set_proxy",
    "set_tenant",
    "set_egress",
    "clear_egress",
    "tenant",
    "egress",
    "build_headers",
    "openai",
    "anthropic",
]


# ─── State (per-task, thread-safe via ContextVar) ───────────────────────────
_proxy_key:  ContextVar[Optional[str]] = ContextVar("_cartie_proxy_key",  default=None)
_proxy_base: ContextVar[str]           = ContextVar("_cartie_proxy_base", default="https://proxy.cartieai.com")
_tenant:     ContextVar[Optional[str]] = ContextVar("_cartie_tenant",     default=None)
_egress:     ContextVar[Optional[dict]] = ContextVar("_cartie_egress",    default=None)


# ─── Public configuration ───────────────────────────────────────────────────
def set_proxy(key: str, base: Optional[str] = None) -> None:
    """Configure the Cartie proxy key + (optional) base URL.

    Args:
        key:  The proxy key generated at POST /api/llm-proxy/keys.
        base: Override the proxy base URL (default https://proxy.cartieai.com).
              Use http://localhost:8001/api/llm-proxy for local development.
    """
    _proxy_key.set(key)
    if base:
        _proxy_base.set(base.rstrip("/"))


def set_tenant(label: str) -> None:
    """Set the X-Cartie-Tenant label for all subsequent LLM calls in this task.

    Args:
        label: A short, lowercase department / cost-center identifier,
               e.g. "marketing", "product", "sales-emea".
    """
    _tenant.set(label.strip())


def set_egress(
    *,
    source: str,
    dest: str,
    data_gb: float,
    token_pct: float = 10.0,
    enforce: bool = False,
) -> None:
    """Set the egress preflight hint that will be sent on every call.

    Args:
        source:    Source cloud — "aws" | "gcp" | "azure" | "snowflake".
        dest:      Destination — "aws" | "gcp" | "azure" | "internet" | "same-region".
        data_gb:   Payload size in gigabytes.
        token_pct: Percent of the payload that becomes LLM input tokens (default 10).
        enforce:   When True, proxy returns 412 if the verdict is rose.
    """
    _egress.set({
        "source": source,
        "dest": dest,
        "data_gb": float(data_gb),
        "token_pct": float(token_pct),
        "enforce": bool(enforce),
    })


def clear_egress() -> None:
    """Remove any egress hint so subsequent calls skip the preflight."""
    _egress.set(None)


# ─── Context-manager scoping ────────────────────────────────────────────────
@contextmanager
def tenant(label: str) -> Iterator[None]:
    """Scope a block of code to a specific tenant label.

    Example:
        with cartie.tenant("sales-emea"):
            client.chat.completions.create(...)
    """
    tok = _tenant.set(label.strip())
    try:
        yield
    finally:
        _tenant.reset(tok)


@contextmanager
def egress(
    *,
    source: str,
    dest: str,
    data_gb: float,
    token_pct: float = 10.0,
    enforce: bool = False,
) -> Iterator[None]:
    """Scope a block of code to a specific egress preflight hint."""
    tok = _egress.set({
        "source": source, "dest": dest, "data_gb": float(data_gb),
        "token_pct": float(token_pct), "enforce": bool(enforce),
    })
    try:
        yield
    finally:
        _egress.reset(tok)


# ─── Header builder (exposed for testing / custom transports) ──────────────
def build_headers() -> dict[str, str]:
    """Returns the Cartie-specific headers to inject on the next call."""
    h: dict[str, str] = {}
    pk = _proxy_key.get()
    if pk:
        h["X-Cartie-Key"] = pk
    t = _tenant.get()
    if t:
        h["X-Cartie-Tenant"] = t
    eg = _egress.get()
    if eg:
        h["X-Cartie-Source-Cloud"] = str(eg["source"])
        h["X-Cartie-Dest-Cloud"] = str(eg["dest"])
        h["X-Cartie-Data-GB"] = str(eg["data_gb"])
        h["X-Cartie-Token-Budget-Pct"] = str(eg["token_pct"])
        if eg.get("enforce"):
            h["X-Cartie-Egress-Enforce"] = "true"
    return h


# ─── Drop-in client factories ───────────────────────────────────────────────
def openai(**kwargs: Any):
    """Returns an `openai.OpenAI` client routed through the Cartie proxy.

    Drop-in replacement for `openai.OpenAI(...)`. Every call auto-carries
    the headers configured via set_proxy / set_tenant / set_egress.

    Requires the `openai` package: `pip install openai`.
    """
    try:
        from openai import OpenAI  # type: ignore
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "Install the openai SDK first: `pip install openai`"
        ) from exc
    base = kwargs.pop("base_url", f"{_proxy_base.get()}/v1")
    headers = {**build_headers(), **kwargs.pop("default_headers", {})}
    return OpenAI(base_url=base, default_headers=headers, **kwargs)


def anthropic(**kwargs: Any):
    """Returns an `anthropic.Anthropic` client routed through the Cartie proxy.

    Drop-in replacement for `anthropic.Anthropic(...)`.

    Requires the `anthropic` package: `pip install anthropic`.
    """
    try:
        from anthropic import Anthropic  # type: ignore
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "Install the anthropic SDK first: `pip install anthropic`"
        ) from exc
    base = kwargs.pop("base_url", f"{_proxy_base.get()}/v1")
    headers = {**build_headers(), **kwargs.pop("default_headers", {})}
    return Anthropic(base_url=base, default_headers=headers, **kwargs)