import os
from typing import Any, Callable
from functools import lru_cache
import time
import re

from smolagents import LiteLLMModel
from litellm import RateLimitError

# IMPORTANT: Set these environment variables BEFORE importing/using litellm
# This tells LiteLLM to route through OpenRouter
"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"
os.environ["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY", )
os.environ["OPENROUTER_API_BASE"] = "https://openrouter.ai/api/v1"


class WrapperLiteLLMModel(LiteLLMModel):
    """Wrapper for LiteLLMModel with retry logic for rate limits."""
    
    def __call__(self, messages, **kwargs):
        max_retry = 5
        for attempt in range(max_retry):
            try:
                return super().__call__(messages, **kwargs)
            except RateLimitError as e:
                print(f"RateLimitError (attempt {attempt+1}/{max_retry})")
                match = re.search(r'"retryDelay": ?"(\d+)s"', str(e))
                retry_seconds = int(match.group(1)) if match else 50
                print(f"Sleeping for {retry_seconds} seconds before retrying...")
                time.sleep(retry_seconds)
            except Exception as e:
                print(f"Error: {type(e).__name__}: {e}")
                raise

        raise RateLimitError(f"Rate limit exceeded after {max_retry} retries.")


def get_lite_llm_model(model_id: str, **kwargs) -> WrapperLiteLLMModel:
    """
    Returns a LiteLLM model instance configured for OpenRouter.

    Args:
        model_id (str): The model identifier in OpenRouter format.
                       Example: "openrouter/google/gemini-2.0-flash-001"
        **kwargs: Additional keyword arguments for the model.

    Returns:
        WrapperLiteLLMModel: LiteLLM model instance.
    """
    api_key = os.getenv("OPENROUTER_API_KEY", "")
    
    if not api_key:
        raise ValueError(
            "OPENROUTER_API_KEY environment variable is not set. "
            "Please set it with: export OPENROUTER_API_KEY='your-key-here'"
        )
    
    # Ensure the model_id starts with "openrouter/"
    if not model_id.startswith("openrouter/"):
        print(f"Warning: model_id should start with 'openrouter/'. Got: {model_id}")
        print(f"Prepending 'openrouter/' to model_id")
        model_id = f"openrouter/{model_id}"
    
    print(f"Initializing LiteLLM model: {model_id}")
    
    return WrapperLiteLLMModel(
        model_id=model_id,
        api_key=api_key,
        # Don't set api_base here - let litellm handle it via OPENROUTER_API_BASE env var
        **kwargs
    )


def get_model(model_type: str, model_id: str, **kwargs) -> Any:
    """
    Returns a model instance based on the specified type.

    Args:
        model_type (str): The type of the model ("LiteLLMModel").
        model_id (str): The model identifier.
        **kwargs: Additional keyword arguments for the model.

    Returns:
        Any: Model instance of the specified type.
    """
    models: dict[str, Callable[..., Any]] = {
        "LiteLLMModel": get_lite_llm_model,
    }

    if model_type not in models:
        raise ValueError(f"Unknown model type: {model_type}")

    return models[model_type](model_id, **kwargs)


# ============================================================================
# AVAILABLE OPENROUTER MODELS (as of 2025)
# ============================================================================
# Format: "openrouter/<provider>/<model-name>"
#
# FREE MODELS (good for testing):
#   - openrouter/google/gemini-2.0-flash-001
#   - openrouter/meta-llama/llama-3.2-3b-instruct:free
#   - openrouter/microsoft/phi-3-mini-128k-instruct:free
#   - openrouter/qwen/qwen-2-7b-instruct:free
#
# PAID MODELS (better quality):
#   - openrouter/anthropic/claude-3.5-sonnet
#   - openrouter/anthropic/claude-3-haiku
#   - openrouter/openai/gpt-4o-mini
#   - openrouter/openai/gpt-4o
#   - openrouter/google/gemini-pro
#   - openrouter/google/gemini-1.5-flash
#   - openrouter/meta-llama/llama-3.1-70b-instruct
#   - openrouter/mistralai/mistral-large
#
# Check https://openrouter.ai/models for the full list and current availability
# ============================================================================


if __name__ == "__main__":
    # Test the model
    print("Testing OpenRouter connection...")
    
    try:
        # Use a free model for testing
        model = get_model("LiteLLMModel", "openrouter/google/gemini-2.0-flash-001")
        
        # Simple test
        messages = [{"role": "user", "content": "Say hello in one word."}]
        response = model(messages)
        print(f"Response: {response}")
        print("✅ OpenRouter connection successful!")
        
    except Exception as e:
        print(f"❌ Error: {type(e).__name__}: {e}")
        print("\nTroubleshooting:")
        print("1. Make sure OPENROUTER_API_KEY is set")
        print("2. Check your API key at https://openrouter.ai/keys")
        print("3. Verify you have credits at https://openrouter.ai/credits")