"""
LiteLLM integration for cascadeflow.

Provides accurate cost tracking using LiteLLM's pricing database,
which is maintained and updated regularly by the LiteLLM team.

NEW in v0.2.0 (Phase 2, Milestone 2.1):
    - LiteLLMCostProvider: Accurate cost calculations using LiteLLM
    - SUPPORTED_PROVIDERS: Strategic provider selection with value props
    - Provider validation
    - Automatic fallback if LiteLLM not installed

Benefits over custom pricing:
    - ✓ Always up-to-date pricing (LiteLLM team maintains it)
    - ✓ Covers 100+ models across 10+ providers
    - ✓ Includes both input and output token pricing
    - ✓ Handles special pricing (batch, cached tokens, etc.)

Usage:
    >>> from cascadeflow.integrations.litellm import LiteLLMCostProvider
    >>>
    >>> # Create cost provider
    >>> cost_provider = LiteLLMCostProvider()
    >>>
    >>> # Calculate cost
    >>> cost = cost_provider.calculate_cost(
    ...     model="gpt-4",
    ...     input_tokens=100,
    ...     output_tokens=50
    ... )
    >>> print(f"Cost: ${cost:.6f}")

Installation:
    pip install litellm

    Optional for even more providers:
    pip install litellm[extra_providers]
"""

import logging
from dataclasses import dataclass
from typing import Optional

logger = logging.getLogger(__name__)

# Try to import LiteLLM (optional dependency)
try:
    import litellm
    from litellm import BudgetManager, completion_cost, model_cost

    LITELLM_AVAILABLE = True
    BUDGET_MANAGER_AVAILABLE = True
    logger.info("LiteLLM integration available (with BudgetManager)")
except ImportError:
    LITELLM_AVAILABLE = False
    BUDGET_MANAGER_AVAILABLE = False
    BudgetManager = None
    logger.warning(
        "LiteLLM not installed. Cost tracking will use fallback estimates. "
        "Install with: pip install litellm"
    )


# ============================================================================
# SUPPORTED PROVIDERS
# ============================================================================


@dataclass
class ProviderInfo:
    """Information about a supported provider."""

    name: str
    display_name: str
    value_prop: str  # Why use this provider?
    pricing_available: bool  # Does LiteLLM have pricing?
    requires_api_key: bool
    example_models: list[str]


# Strategic provider selection (10 providers as per plan)
# Each provider has a clear value proposition
SUPPORTED_PROVIDERS = {
    "openai": ProviderInfo(
        name="openai",
        display_name="OpenAI",
        value_prop="Industry-leading quality, most reliable, best for production",
        pricing_available=True,
        requires_api_key=True,
        example_models=["gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"],
    ),
    "anthropic": ProviderInfo(
        name="anthropic",
        display_name="Anthropic Claude",
        value_prop="Best for reasoning and analysis, strong safety features",
        pricing_available=True,
        requires_api_key=True,
        example_models=[
            "anthropic/claude-opus-4-6-20250610",
            "anthropic/claude-opus-4-5-20251101",
            "anthropic/claude-sonnet-4-5-20250929",
            "anthropic/claude-haiku-4-5-20251001",
        ],
    ),
    "groq": ProviderInfo(
        name="groq",
        display_name="Groq",
        value_prop="Fastest inference speed, ultra-low latency, free tier",
        pricing_available=True,
        requires_api_key=True,
        example_models=[
            "groq/llama-3.1-70b-versatile",
            "groq/llama-3.1-8b-instant",
            "groq/mixtral-8x7b-32768",
        ],
    ),
    "together": ProviderInfo(
        name="together",
        display_name="Together AI",
        value_prop="Cost-effective, wide model selection, good for experimentation",
        pricing_available=True,
        requires_api_key=True,
        example_models=[
            "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
            "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
            "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo",
        ],
    ),
    "huggingface": ProviderInfo(
        name="huggingface",
        display_name="Hugging Face",
        value_prop="Open-source models, community-driven, flexible deployment",
        pricing_available=True,
        requires_api_key=True,
        example_models=[
            "huggingface/mistralai/Mistral-7B-Instruct-v0.2",
            "huggingface/meta-llama/Llama-2-70b-chat",
        ],
    ),
    "ollama": ProviderInfo(
        name="ollama",
        display_name="Ollama",
        value_prop="Local/on-prem deployment, zero cost, full privacy",
        pricing_available=False,  # Free, local
        requires_api_key=False,
        example_models=["llama3.1:8b", "llama3.1:70b", "mistral", "codellama"],
    ),
    "vllm": ProviderInfo(
        name="vllm",
        display_name="vLLM",
        value_prop="Self-hosted inference, high throughput, production-ready",
        pricing_available=False,  # Self-hosted
        requires_api_key=False,
        example_models=["meta-llama/Llama-3.1-70B", "meta-llama/Llama-3.1-8B"],
    ),
    "google": ProviderInfo(
        name="google",
        display_name="Google (Vertex AI)",
        value_prop="Enterprise integration, GCP ecosystem, Gemini models",
        pricing_available=True,
        requires_api_key=True,
        example_models=["gemini/gemini-pro", "gemini/gemini-1.5-pro", "gemini/gemini-1.5-flash"],
    ),
    "azure": ProviderInfo(
        name="azure",
        display_name="Azure OpenAI",
        value_prop="Enterprise compliance, HIPAA/SOC2, Microsoft ecosystem",
        pricing_available=True,
        requires_api_key=True,
        example_models=["azure/gpt-4", "azure/gpt-4-turbo", "azure/gpt-3.5-turbo"],
    ),
    "deepseek": ProviderInfo(
        name="deepseek",
        display_name="DeepSeek",
        value_prop="Specialized code models, very cost-effective for coding tasks",
        pricing_available=True,
        requires_api_key=True,
        example_models=["deepseek/deepseek-coder", "deepseek/deepseek-chat"],
    ),
}


def validate_provider(provider: str) -> bool:
    """
    Validate if provider is supported.

    Args:
        provider: Provider name to validate

    Returns:
        True if supported, False otherwise

    Example:
        >>> validate_provider("openai")
        True
        >>> validate_provider("unknown_provider")
        False
    """
    supported = provider.lower() in SUPPORTED_PROVIDERS

    if not supported:
        available = ", ".join(SUPPORTED_PROVIDERS.keys())
        logger.warning(f"Provider '{provider}' not in supported list. " f"Available: {available}")

    return supported


def get_provider_info(provider: str) -> Optional[ProviderInfo]:
    """
    Get information about a provider.

    Args:
        provider: Provider name

    Returns:
        ProviderInfo if found, None otherwise

    Example:
        >>> info = get_provider_info("groq")
        >>> print(info.value_prop)
        'Fastest inference speed, ultra-low latency, free tier'
    """
    return SUPPORTED_PROVIDERS.get(provider.lower())


# ============================================================================
# LITELLM COST PROVIDER
# ============================================================================


class LiteLLMCostProvider:
    """
    Cost calculation using LiteLLM's pricing database.

    This provides accurate, up-to-date pricing for 100+ models across
    10+ providers without maintaining custom pricing tables.

    Example:
        >>> cost_provider = LiteLLMCostProvider()
        >>>
        >>> # Calculate cost from token counts
        >>> cost = cost_provider.calculate_cost(
        ...     model="gpt-4",
        ...     input_tokens=100,
        ...     output_tokens=50
        ... )
        >>>
        >>> # Get model pricing info
        >>> pricing = cost_provider.get_model_cost("gpt-4")
        >>> print(f"Input: ${pricing['input_cost_per_token']:.8f}/token")
    """

    def __init__(self, fallback_enabled: bool = True):
        """
        Initialize LiteLLM cost provider.

        Args:
            fallback_enabled: Use fallback estimates if LiteLLM unavailable
        """
        self.fallback_enabled = fallback_enabled

        if not LITELLM_AVAILABLE:
            logger.warning("LiteLLM not available. Cost calculations will use fallback estimates.")

    def calculate_cost(
        self,
        model: str,
        input_tokens: int = 0,
        output_tokens: int = 0,
        **kwargs,
    ) -> float:
        """
        Calculate cost using LiteLLM.

        Args:
            model: Model name (e.g., "gpt-4", "claude-3-opus")
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            **kwargs: Additional LiteLLM parameters (cache_tokens, etc.)

        Returns:
            Cost in USD

        Example:
            >>> cost = provider.calculate_cost("gpt-4", 100, 50)
            >>> print(f"${cost:.6f}")
            $0.004500
        """
        override_pricing = {
            "claude-opus-4-6": {"input": 5.0, "output": 25.0},
            "claude-opus-4-5": {"input": 5.0, "output": 25.0},
            "claude-opus-4": {"input": 15.0, "output": 75.0},
            "claude-sonnet-4-5": {"input": 3.0, "output": 15.0},
            "claude-haiku-4-5": {"input": 1.0, "output": 5.0},
            "claude-haiku-4.5": {"input": 1.0, "output": 5.0},
            "claude-3-5-haiku": {"input": 1.0, "output": 5.0},
        }
        for model_prefix, rates in override_pricing.items():
            if model.startswith(model_prefix):
                input_cost = (input_tokens / 1_000_000) * rates["input"]
                output_cost = (output_tokens / 1_000_000) * rates["output"]
                return input_cost + output_cost

        if not LITELLM_AVAILABLE:
            if self.fallback_enabled:
                return self._fallback_cost(model, input_tokens, output_tokens)
            else:
                raise RuntimeError(
                    "LiteLLM not installed and fallback disabled. "
                    "Install with: pip install litellm"
                )

        try:
            # Create a mock completion response object for LiteLLM
            # LiteLLM's API expects a response object with usage information
            from litellm import ModelResponse

            mock_response = ModelResponse(
                id="mock",
                model=model,
                choices=[{"message": {"content": ""}, "finish_reason": "stop"}],
                usage={
                    "prompt_tokens": input_tokens,
                    "completion_tokens": output_tokens,
                    "total_tokens": input_tokens + output_tokens,
                },
            )

            # Calculate cost using the mock response
            cost = completion_cost(
                completion_response=mock_response,
                model=model,
                **kwargs,
            )

            logger.debug(
                f"LiteLLM cost for {model}: ${cost:.6f} "
                f"({input_tokens} in, {output_tokens} out)"
            )

            return cost

        except Exception as e:
            logger.warning(f"LiteLLM cost calculation failed for {model}: {e}")

            if self.fallback_enabled:
                return self._fallback_cost(model, input_tokens, output_tokens)
            else:
                raise

    def get_model_cost(self, model: str) -> dict:
        """
        Get pricing information for a model.

        Args:
            model: Model name

        Returns:
            Dict with pricing info:
                - input_cost_per_token: Cost per input token (USD)
                - output_cost_per_token: Cost per output token (USD)
                - max_tokens: Maximum context length
                - supports_streaming: Whether streaming is supported

        Example:
            >>> pricing = provider.get_model_cost("gpt-4")
            >>> print(f"Input: ${pricing['input_cost_per_token']:.8f}/token")
            Input: $0.00003000/token
        """
        if not LITELLM_AVAILABLE:
            logger.warning("LiteLLM not available, returning fallback pricing")
            return self._fallback_pricing(model)

        try:
            # Get pricing from LiteLLM's model_cost dict
            pricing = model_cost.get(model, {})

            if not pricing:
                # Try using completion_cost to derive pricing (handles provider prefixes better)
                try:
                    from litellm import ModelResponse

                    # Test with 1M tokens to get per-token costs
                    mock_response = ModelResponse(
                        id="mock",
                        model=model,
                        choices=[{"message": {"content": ""}, "finish_reason": "stop"}],
                        usage={
                            "prompt_tokens": 1_000_000,
                            "completion_tokens": 1_000_000,
                            "total_tokens": 2_000_000,
                        },
                    )

                    total_cost = completion_cost(completion_response=mock_response, model=model)

                    # Derive per-token costs (assuming equal input/output in the mock)
                    # completion_cost will use actual model pricing
                    total_cost / 2_000_000

                    # Get separate costs by testing with just input tokens
                    mock_input_only = ModelResponse(
                        id="mock",
                        model=model,
                        choices=[{"message": {"content": ""}, "finish_reason": "stop"}],
                        usage={
                            "prompt_tokens": 1_000_000,
                            "completion_tokens": 0,
                            "total_tokens": 1_000_000,
                        },
                    )
                    input_cost = completion_cost(completion_response=mock_input_only, model=model)
                    input_cost_per_token = input_cost / 1_000_000

                    # Calculate output cost per token
                    output_cost = total_cost - input_cost
                    output_cost_per_token = output_cost / 1_000_000

                    logger.debug(f"Derived pricing for {model} using completion_cost")

                    return {
                        "input_cost_per_token": input_cost_per_token,
                        "output_cost_per_token": output_cost_per_token,
                        "max_tokens": 4096,  # Default
                        "supports_streaming": True,
                    }

                except Exception as e:
                    logger.debug(f"Could not derive pricing for {model}: {e}")
                    return self._fallback_pricing(model)

            return {
                "input_cost_per_token": pricing.get("input_cost_per_token", 0),
                "output_cost_per_token": pricing.get("output_cost_per_token", 0),
                "max_tokens": pricing.get("max_tokens", 4096),
                "supports_streaming": pricing.get("supports_streaming", True),
            }

        except Exception as e:
            logger.warning(f"Error getting pricing for {model}: {e}")
            return self._fallback_pricing(model)

    def _fallback_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """
        Fallback cost estimation when LiteLLM unavailable.

        Uses rough estimates based on typical pricing.
        """
        # Rough pricing estimates (per 1M tokens)
        rough_pricing = {
            # OpenAI
            "gpt-4": {"input": 30.0, "output": 60.0},
            "gpt-4-turbo": {"input": 10.0, "output": 30.0},
            "gpt-4o": {"input": 5.0, "output": 15.0},
            "gpt-4o-mini": {"input": 0.15, "output": 0.60},
            "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
            # Anthropic 4.x
            "claude-opus-4-6": {"input": 5.0, "output": 25.0},
            "claude-opus-4-5": {"input": 5.0, "output": 25.0},
            "claude-sonnet-4-5": {"input": 3.0, "output": 15.0},
            "claude-haiku-4-5": {"input": 1.0, "output": 5.0},
            # Anthropic 3.x
            "claude-3-opus": {"input": 15.0, "output": 75.0},
            "claude-3-5-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-sonnet": {"input": 3.0, "output": 15.0},
            "claude-3-haiku": {"input": 0.25, "output": 1.25},
            # Default
            "default": {"input": 1.0, "output": 2.0},
        }

        # Get pricing by exact match first, then prefix match, then default
        pricing = rough_pricing.get(model)
        if pricing is None:
            for model_prefix, prefix_pricing in rough_pricing.items():
                if model_prefix != "default" and model.startswith(model_prefix):
                    pricing = prefix_pricing
                    break
        if pricing is None:
            pricing = rough_pricing["default"]

        input_cost = (input_tokens / 1_000_000) * pricing["input"]
        output_cost = (output_tokens / 1_000_000) * pricing["output"]

        total_cost = input_cost + output_cost

        logger.debug(
            f"Fallback cost for {model}: ${total_cost:.6f} "
            f"({input_tokens} in @ ${pricing['input']}/1M, "
            f"{output_tokens} out @ ${pricing['output']}/1M)"
        )

        return total_cost

    def _fallback_pricing(self, model: str) -> dict:
        """Fallback pricing info when LiteLLM unavailable."""
        return {
            "input_cost_per_token": 0.000001,  # $1/1M tokens
            "output_cost_per_token": 0.000002,  # $2/1M tokens
            "max_tokens": 4096,
            "supports_streaming": True,
        }


# ============================================================================
# CONVENIENCE FUNCTIONS
# ============================================================================


def get_model_cost(model: str) -> dict:
    """
    Get pricing information for a model.

    Convenience function that creates a LiteLLMCostProvider and calls get_model_cost().

    Args:
        model: Model name

    Returns:
        Dict with pricing info

    Example:
        >>> pricing = get_model_cost("gpt-4")
        >>> print(f"Input: ${pricing['input_cost_per_token']:.8f}/token")
    """
    provider = LiteLLMCostProvider()
    return provider.get_model_cost(model)


def calculate_cost(
    model: str,
    input_tokens: int = 0,
    output_tokens: int = 0,
    **kwargs,
) -> float:
    """
    Calculate cost for a model call.

    Convenience function that creates a LiteLLMCostProvider and calls calculate_cost().

    Args:
        model: Model name
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        **kwargs: Additional parameters

    Returns:
        Cost in USD

    Example:
        >>> cost = calculate_cost("gpt-4", input_tokens=100, output_tokens=50)
        >>> print(f"${cost:.6f}")
        $0.004500
    """
    provider = LiteLLMCostProvider()
    return provider.calculate_cost(model, input_tokens, output_tokens, **kwargs)


# ============================================================================
# LITELLM BUDGET TRACKER
# ============================================================================


class LiteLLMBudgetTracker:
    """
    Budget tracking using LiteLLM's BudgetManager.

    This integrates with LiteLLM's BudgetManager for actual spending tracking,
    not just cost estimation. Tracks real API call costs per user.

    NEW in v0.2.0 (Phase 2, Milestone 2.1 - Enhanced):
        - Integration with LiteLLM's BudgetManager
        - Actual spending tracking (not just estimates)
        - Per-user budgets with real-time enforcement
        - Compatible with cascadeflow's CostTracker

    Example:
        >>> tracker = LiteLLMBudgetTracker()
        >>>
        >>> # Set user budget
        >>> tracker.set_user_budget("user_123", max_budget=10.0)
        >>>
        >>> # Update cost after API call
        >>> tracker.update_cost(
        ...     user="user_123",
        ...     model="gpt-4",
        ...     prompt_tokens=100,
        ...     completion_tokens=50
        ... )
        >>>
        >>> # Check if user can afford more
        >>> can_continue = tracker.can_user_afford("user_123", estimated_cost=0.05)
    """

    def __init__(self, fallback_to_cascadeflow: bool = True):
        """
        Initialize LiteLLM budget tracker.

        Args:
            fallback_to_cascadeflow: Use cascadeflow's CostTracker if LiteLLM unavailable
        """
        self.fallback_to_cascadeflow = fallback_to_cascadeflow
        self.budget_manager = None
        self.cost_provider = LiteLLMCostProvider()
        self._user_budgets: dict[str, dict] = {}

        if BUDGET_MANAGER_AVAILABLE:
            self.budget_manager = BudgetManager(project_name="cascadeflow")
            logger.info("LiteLLM BudgetManager initialized")
        else:
            logger.warning(
                "LiteLLM BudgetManager not available. "
                "Install litellm for actual spending tracking."
            )

            if fallback_to_cascadeflow:
                # Import cascadeflow's CostTracker as fallback
                try:
                    from cascadeflow.telemetry import CostTracker

                    self.cost_tracker = CostTracker()
                    logger.info("Using cascadeflow CostTracker as fallback")
                except ImportError:
                    self.cost_tracker = None
                    logger.warning("cascadeflow CostTracker also unavailable")

    def set_user_budget(self, user: str, max_budget: float) -> None:
        """
        Set maximum budget for a user.

        Args:
            user: User identifier
            max_budget: Maximum budget in USD

        Example:
            >>> tracker.set_user_budget("user_123", max_budget=10.0)
        """
        self._user_budgets[user] = {
            "max_budget": max_budget,
            "current_cost": 0.0,
        }

        if self.budget_manager:
            try:
                self.budget_manager.create_budget(user=user, total_budget=max_budget)
            except Exception as e:
                logger.debug(f"BudgetManager.create_budget failed for {user}: {e}")

        logger.info(f"Set budget for {user}: ${max_budget:.2f}")

    def update_cost(
        self,
        user: str,
        model: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        response: Optional[dict] = None,
    ) -> float:
        """
        Update cost after API call.

        Can use either token counts or actual API response.

        Args:
            user: User identifier
            model: Model name
            prompt_tokens: Input tokens used
            completion_tokens: Output tokens used
            response: Optional API response dict (for more accurate tracking)

        Returns:
            Cost of this call in USD

        Example:
            >>> # From token counts
            >>> cost = tracker.update_cost(
            ...     user="user_123",
            ...     model="gpt-4",
            ...     prompt_tokens=100,
            ...     completion_tokens=50
            ... )
            >>>
            >>> # From API response
            >>> cost = tracker.update_cost(
            ...     user="user_123",
            ...     response=api_response
            ... )
        """
        # Calculate cost from tokens or response
        if response:
            cost = self.cost_provider.calculate_cost(
                model=model,
                input_tokens=prompt_tokens,
                output_tokens=completion_tokens,
            )
        else:
            cost = self.cost_provider.calculate_cost(
                model=model,
                input_tokens=prompt_tokens,
                output_tokens=completion_tokens,
            )

        # Track in internal budget dict
        if user in self._user_budgets:
            self._user_budgets[user]["current_cost"] += cost

        # Also track in cascadeflow CostTracker if available
        if self.fallback_to_cascadeflow and hasattr(self, "cost_tracker") and self.cost_tracker:
            self.cost_tracker.add_cost(
                model=model,
                provider="",
                tokens=prompt_tokens + completion_tokens,
                cost=cost,
                user_id=user,
            )

        logger.debug(f"Updated cost for {user}: ${cost:.6f}")
        return cost

    def get_user_budget(self, user: str) -> dict:
        """
        Get user's budget information.

        Args:
            user: User identifier

        Returns:
            Dict with budget info:
                - max_budget: Maximum budget (USD)
                - current_cost: Current spending (USD)
                - remaining: Remaining budget (USD)
                - exceeded: Whether budget exceeded

        Example:
            >>> info = tracker.get_user_budget("user_123")
            >>> print(f"Spent: ${info['current_cost']:.2f}")
            >>> print(f"Remaining: ${info['remaining']:.2f}")
        """
        budget = self._user_budgets.get(user)
        if budget:
            max_budget = budget["max_budget"]
            current_cost = budget["current_cost"]
            remaining = max_budget - current_cost
            exceeded = current_cost > max_budget

            return {
                "max_budget": max_budget,
                "current_cost": current_cost,
                "remaining": remaining,
                "exceeded": exceeded,
            }

        return {
            "max_budget": 0,
            "current_cost": 0,
            "remaining": 0,
            "exceeded": False,
        }

    def can_user_afford(self, user: str, estimated_cost: float) -> bool:
        """
        Check if user can afford estimated cost.

        Args:
            user: User identifier
            estimated_cost: Estimated cost in USD

        Returns:
            True if user can afford, False otherwise

        Example:
            >>> if tracker.can_user_afford("user_123", 0.05):
            ...     # Make API call
            ...     pass
            ... else:
            ...     # User over budget
            ...     print("Budget exceeded")
        """
        budget_info = self.get_user_budget(user)

        # If already exceeded, can't afford
        if budget_info["exceeded"]:
            return False

        # Check if remaining budget covers estimated cost
        return budget_info["remaining"] >= estimated_cost

    def reset_user_budget(self, user: str) -> None:
        """
        Reset user's spending (budget remains same).

        Args:
            user: User identifier

        Example:
            >>> tracker.reset_user_budget("user_123")
        """
        if user in self._user_budgets:
            self._user_budgets[user]["current_cost"] = 0.0
            logger.info(f"Reset budget for {user}")


# ============================================================================
# LITELLM CALLBACKS FOR CASCADEFLOW INTEGRATION
# ============================================================================


class cascadeflowLiteLLMCallback:
    """
    Custom LiteLLM callback for cascadeflow integration.

    Bridges LiteLLM callbacks to cascadeflow's telemetry systems:
    - CostTracker for spending tracking
    - MetricsCollector for analytics
    - CallbackManager for custom handlers

    NEW in v0.2.0 (Phase 2, Milestone 2.1 - Enhanced):
        - Automatic cost tracking with cascadeflow systems
        - Success/failure callbacks
        - Integration with existing telemetry

    Example:
        >>> # Set up callback
        >>> callback = cascadeflowLiteLLMCallback()
        >>>
        >>> # Register with LiteLLM (if installed)
        >>> if LITELLM_AVAILABLE:
        ...     import litellm
        ...     litellm.success_callback = [callback.log_success]
        ...     litellm.failure_callback = [callback.log_failure]
        >>>
        >>> # Now all LiteLLM calls automatically tracked in cascadeflow!
    """

    def __init__(
        self,
        cost_tracker=None,
        metrics_collector=None,
        callback_manager=None,
    ):
        """
        Initialize LiteLLM callback for cascadeflow.

        Args:
            cost_tracker: cascadeflow CostTracker instance (optional)
            metrics_collector: cascadeflow MetricsCollector instance (optional)
            callback_manager: cascadeflow CallbackManager instance (optional)
        """
        self.cost_tracker = cost_tracker
        self.metrics_collector = metrics_collector
        self.callback_manager = callback_manager

        # Try to import cascadeflow components if not provided
        if cost_tracker is None:
            try:
                from cascadeflow.telemetry import CostTracker

                self.cost_tracker = CostTracker()
                logger.info("cascadeflow CostTracker initialized for LiteLLM callbacks")
            except ImportError:
                logger.debug("CostTracker not available for callbacks")

        if metrics_collector is None:
            try:
                from cascadeflow.telemetry import MetricsCollector

                self.metrics_collector = MetricsCollector()
                logger.info("MetricsCollector initialized for LiteLLM callbacks")
            except ImportError:
                logger.debug("MetricsCollector not available for callbacks")

    def log_success(self, kwargs, response_obj, start_time, end_time):
        """
        Called after successful LiteLLM API call.

        Tracks costs and metrics automatically, including logprobs if available.

        Args:
            kwargs: Request parameters
            response_obj: API response
            start_time: Request start time
            end_time: Request end time
        """
        try:
            # Extract info from response
            model = kwargs.get("model", "unknown")
            user_id = kwargs.get("user", None)

            # Get cost from response if available
            cost = kwargs.get("response_cost", 0)

            # Get token counts
            prompt_tokens = getattr(response_obj, "prompt_tokens", 0)
            completion_tokens = getattr(response_obj, "completion_tokens", 0)
            total_tokens = prompt_tokens + completion_tokens

            # Calculate latency
            latency_ms = (end_time - start_time) * 1000

            # Extract logprobs if available (NEW)
            logprobs_data = None
            has_logprobs = False
            avg_logprob = None

            try:
                # LiteLLM returns logprobs in response_obj.choices[0].logprobs
                if hasattr(response_obj, "choices") and len(response_obj.choices) > 0:
                    choice = response_obj.choices[0]
                    if hasattr(choice, "logprobs") and choice.logprobs:
                        has_logprobs = True
                        logprobs_data = choice.logprobs

                        # Extract content logprobs for confidence scoring
                        if hasattr(logprobs_data, "content") and logprobs_data.content:
                            logprob_values = [
                                item.logprob
                                for item in logprobs_data.content
                                if hasattr(item, "logprob")
                            ]
                            if logprob_values:
                                avg_logprob = sum(logprob_values) / len(logprob_values)

                                logger.debug(
                                    f"LiteLLM logprobs extracted: {len(logprob_values)} tokens, "
                                    f"avg logprob: {avg_logprob:.3f}"
                                )
            except Exception as e:
                logger.debug(f"Could not extract logprobs from LiteLLM response: {e}")

            # Build metadata with logprobs info
            metadata = {
                "latency_ms": latency_ms,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "has_logprobs": has_logprobs,
            }

            if avg_logprob is not None:
                metadata["avg_logprob"] = avg_logprob
                # Convert logprob to confidence (e^logprob)
                import math

                metadata["confidence"] = min(0.99, max(0.01, math.exp(avg_logprob)))

            # Track cost if tracker available
            if self.cost_tracker and cost > 0:
                self.cost_tracker.add_cost(
                    model=model,
                    provider=kwargs.get("custom_llm_provider", ""),
                    tokens=total_tokens,
                    cost=cost,
                    user_id=user_id,
                    metadata=metadata,
                )

                log_msg = (
                    f"LiteLLM success tracked: {model}, ${cost:.6f}, "
                    f"{total_tokens} tokens, {latency_ms:.0f}ms"
                )
                if has_logprobs and avg_logprob is not None:
                    log_msg += f", confidence: {metadata['confidence']:.2f}"

                logger.debug(log_msg)

            # Track metrics if collector available
            if self.metrics_collector:
                # Would integrate with MetricsCollector here
                # Could use logprobs for quality metrics
                pass

            # Trigger custom callbacks if manager available
            if self.callback_manager:
                # Would trigger custom handlers here
                pass

        except Exception as e:
            logger.error(f"Error in LiteLLM success callback: {e}")

    def log_failure(self, kwargs, response_obj, start_time, end_time):
        """
        Called after failed LiteLLM API call.

        Logs failures for monitoring.

        Args:
            kwargs: Request parameters
            response_obj: Error response
            start_time: Request start time
            end_time: Request end time
        """
        try:
            model = kwargs.get("model", "unknown")
            user_id = kwargs.get("user", None)
            error = str(response_obj)

            logger.warning(f"LiteLLM call failed: model={model}, user={user_id}, " f"error={error}")

            # Could integrate with error tracking systems here
            if self.callback_manager:
                # Would trigger error handlers here
                pass

        except Exception as e:
            logger.error(f"Error in LiteLLM failure callback: {e}")


def setup_litellm_callbacks(
    cost_tracker=None,
    metrics_collector=None,
    callback_manager=None,
) -> bool:
    """
    Set up LiteLLM callbacks for automatic cascadeflow integration.

    Call this once at startup to enable automatic tracking of all
    LiteLLM calls in your cascadeflow telemetry systems.

    Args:
        cost_tracker: cascadeflow CostTracker instance (optional)
        metrics_collector: cascadeflow MetricsCollector instance (optional)
        callback_manager: cascadeflow CallbackManager instance (optional)

    Returns:
        True if callbacks set up successfully, False if LiteLLM not available

    Example:
        >>> from cascadeflow.integrations.litellm import setup_litellm_callbacks
        >>> from cascadeflow.telemetry import CostTracker
        >>>
        >>> # Set up tracking
        >>> tracker = CostTracker()
        >>> setup_litellm_callbacks(cost_tracker=tracker)
        >>>
        >>> # Now all LiteLLM calls automatically tracked!
        >>> import litellm
        >>> response = litellm.completion(
        ...     model="gpt-4",
        ...     messages=[{"role": "user", "content": "Hello"}]
        ... )
        >>> # Cost automatically added to tracker ✓
    """
    if not LITELLM_AVAILABLE:
        logger.warning(
            "LiteLLM not installed. Cannot set up callbacks. " "Install with: pip install litellm"
        )
        return False

    try:
        import litellm

        # Create callback instance
        callback = cascadeflowLiteLLMCallback(
            cost_tracker=cost_tracker,
            metrics_collector=metrics_collector,
            callback_manager=callback_manager,
        )

        # Register callbacks with LiteLLM
        litellm.success_callback = [callback.log_success]
        litellm.failure_callback = [callback.log_failure]

        logger.info("LiteLLM callbacks registered with cascadeflow telemetry")
        return True

    except Exception as e:
        logger.error(f"Error setting up LiteLLM callbacks: {e}")
        return False


__all__ = [
    "SUPPORTED_PROVIDERS",
    "ProviderInfo",
    "LiteLLMCostProvider",
    "LiteLLMBudgetTracker",
    "cascadeflowLiteLLMCallback",
    "setup_litellm_callbacks",
    "get_model_cost",
    "calculate_cost",
    "validate_provider",
    "get_provider_info",
]
