"""Fenra: Privacy-first AI cost tracking library."""

import os
from typing import Any

from fenra._context import clear_context, get_context, set_context
from fenra._core import enqueue_transaction, flush, init, shutdown

__version__ = "0.1.0"

# Check environment variable for disabling auto-tracking
_auto_track_disabled = os.environ.get("FENRA_DISABLE_AUTO_TRACK", "").lower() in (
    "1",
    "true",
    "yes",
)


def disable_auto_track() -> None:
    """
    Disable auto-tracking/patching of AI provider SDKs.

    This function must be called BEFORE importing fenra if you want to prevent
    automatic patching. Alternatively, set the environment variable
    FENRA_DISABLE_AUTO_TRACK=1 before importing fenra.

    When auto-tracking is disabled, you can still use fenra.track() for manual
    tracking.
    """
    global _auto_track_disabled
    _auto_track_disabled = True


def track(
    provider: str,
    model: str,
    input_tokens: int,
    output_tokens: int,
    *,
    total_tokens: int | None = None,
    reasoning_tokens: int | None = None,
    cached_tokens: int | None = None,
    **context_overrides: Any,
) -> None:
    """
    Manually track an AI call.

    Merges context_overrides with the current context from set_context().

    Args:
        provider: Provider name (e.g., "openai", "anthropic", "custom")
        model: Model identifier (e.g., "gpt-4o", "claude-3-opus")
        input_tokens: Number of input/prompt tokens
        output_tokens: Number of output/completion tokens
        total_tokens: Total tokens (calculated if not provided)
        reasoning_tokens: Reasoning tokens (for o1/o3 models)
        cached_tokens: Cached tokens (if applicable)
        **context_overrides: Additional context to merge with current context
    """
    from fenra._context import get_context

    # Merge context_overrides with current context
    context = get_context()
    context.update(context_overrides)

    # Calculate total_tokens if not provided
    if total_tokens is None:
        total_tokens = input_tokens + output_tokens

    # Build metrics
    metrics: dict[str, Any] = {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens,
    }

    if reasoning_tokens is not None:
        metrics["reasoning_tokens"] = reasoning_tokens

    if cached_tokens is not None:
        metrics["cached_tokens"] = cached_tokens

    # Build transaction
    transaction = {
        "provider": provider,
        "model": model,
        "usage": [
            {
                "type": "tokens",
                "metrics": metrics,
            }
        ],
        "context": context,
    }

    enqueue_transaction(transaction)


# Auto-patch OpenAI when imported (if OpenAI is installed)
def _auto_patch_openai() -> None:
    """Automatically patch OpenAI if it's installed."""
    if _auto_track_disabled:
        return
    try:
        from fenra.integrations.openai import (
            patch_openai,
            patch_openai_async,
            patch_openai_beta_parse,
            patch_openai_beta_parse_async,
            patch_openai_images,
            patch_openai_images_async,
            patch_openai_responses,
            patch_openai_responses_async,
        )

        patch_openai()
        patch_openai_async()
        patch_openai_responses()
        patch_openai_responses_async()
        patch_openai_images()
        patch_openai_images_async()
        patch_openai_beta_parse()
        patch_openai_beta_parse_async()
    except ImportError:
        # OpenAI not installed, that's fine
        pass


# Auto-patch Gemini when imported (if Gemini is installed)
def _auto_patch_gemini() -> None:
    """Automatically patch Gemini if it's installed."""
    if _auto_track_disabled:
        return
    try:
        from fenra.integrations.gemini import (
            patch_gemini,
            patch_gemini_async,
            patch_gemini_stream,
            patch_gemini_stream_async,
        )

        patch_gemini()
        patch_gemini_async()
        patch_gemini_stream()
        patch_gemini_stream_async()
    except ImportError:
        # Gemini not installed, that's fine
        pass


# Auto-patch Anthropic when imported (if Anthropic is installed)
def _auto_patch_anthropic() -> None:
    """Automatically patch Anthropic if it's installed."""
    if _auto_track_disabled:
        return
    try:
        from fenra.integrations.anthropic import (
            patch_anthropic,
            patch_anthropic_async,
            patch_anthropic_stream,
            patch_anthropic_stream_async,
        )

        patch_anthropic()
        patch_anthropic_async()
        patch_anthropic_stream()
        patch_anthropic_stream_async()
    except ImportError:
        # Anthropic not installed, that's fine
        pass


# Auto-patch on import
_auto_patch_openai()
_auto_patch_gemini()
_auto_patch_anthropic()

__all__ = [
    "init",
    "shutdown",
    "flush",
    "track",
    "set_context",  # Primary context API - merges into existing
    "get_context",
    "clear_context",
    "disable_auto_track",
]
