import os

from dotenv import load_dotenv
from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel

load_dotenv()

FREE_MODEL_FALLBACKS = (
    ("OPENROUTER_API_KEY", "openrouter/meta-llama/llama-3.3-70b-instruct:free"),
    ("GROQ_API_KEY", "groq/llama-3.3-70b-versatile"),
    ("GEMINI_API_KEY", "gemini/gemini-2.5-flash"),
)


def looks_like_api_key(value):
    if not value:
        return False
    return value.startswith(("AIza", "sk-", "sk-or-v1-", "gsk_"))


def resolve_api_key(model_id, env=None):
    env = os.environ if env is None else env
    provider_prefix = model_id.split("/", 1)[0].lower()
    provider_key_map = {
        "openrouter": "OPENROUTER_API_KEY",
        "gemini": "GEMINI_API_KEY",
        "groq": "GROQ_API_KEY",
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "together_ai": "TOGETHER_API_KEY",
        "together": "TOGETHER_API_KEY",
    }

    env_var_name = provider_key_map.get(provider_prefix)
    if env_var_name:
        return env.get(env_var_name)

    return env.get("LITELLM_API_KEY") or env.get("API_KEY")


def resolve_model_config(env=None):
    env = os.environ if env is None else env
    model_id = env.get("MODEL_ID")

    if model_id:
        if looks_like_api_key(model_id):
            raise ValueError(
                "MODEL_ID is set to what looks like an API key. "
                "Set MODEL_ID to a provider/model string like 'gemini/gemini-2.5-flash' "
                "and keep the secret in the matching *_API_KEY variable."
            )
        return model_id, resolve_api_key(model_id, env)

    for api_key_name, fallback_model in FREE_MODEL_FALLBACKS:
        api_key = env.get(api_key_name)
        if api_key:
            return fallback_model, api_key

    return FREE_MODEL_FALLBACKS[-1][1], None


def create_agent(env=None):
    model_id, api_key = resolve_model_config(env)
    model = LiteLLMModel(
        model_id=model_id,
        temperature=0.2,
        api_key=api_key,
    )
    search_tool = DuckDuckGoSearchTool()

    return CodeAgent(
        tools=[search_tool],
        model=model,
        additional_authorized_imports=[
            "requests",
            "bs4",
            "json",
            "pandas",
            "numpy",
            "math",
            "datetime",
            "re",
            "csv",
            "statistics",
            "pathlib",
        ],
        max_steps=25,
    )


agent = create_agent()