import os
import datetime as dt
import logging

from smolagents import OpenAIModel, LiteLLMModel, CodeAgent, tool
#
# Lite is for local inference
# openaimodel is for openrouter
#
from smolagents import DuckDuckGoSearchTool, WebSearchTool

from smol_common import setup_logger, log_event, get_model_idshort

import pandas, numpy, sympy
import matplotlib, matplotlib.pyplot

logger = setup_logger(name="smol_codeagent")

MODELS_FREE = [
    "google/gemma-3-27b-it:free",
    "xiaomi/mimo-v2-flash:free",
    "nex-agi/deepseek-v3.1-nex-n1:free"
]

# Google Gemini models with known limitations in CodeAgent
# These may fail with "cookie auth" errors as they require special Google authentication
GEMINI_MODELS = [
    "google/gemini-2.5-flash",
    "google/gemini-2.5-pro",
    "google/gemini-2.0-flash",
    "google/gemini-pro",
    "google/gemini-pro-vision"
]


def run_codeagent(model_id: str, prompt: str, max_steps: int = 10, provider: str = "openrouter"):
    """
    Run vanilla CodeAgent with built-in prompt

    Args:
        model_id: Model identifier to use
        prompt: User prompt/task description
        max_steps: Maximum number of steps (default: 10)
        provider: Provider to use ('openrouter' or 'ollama')

    Returns:
        Result from CodeAgent execution (str or RunResult)

    Raises:
        ValueError: If required API key is not set
        RuntimeError: If model initialization or execution fails
    """
    # Check for API key before starting
    if provider == "openrouter" and not os.environ.get("OPENROUTER_API_KEY"):
        raise ValueError(
            "OPENROUTER_API_KEY environment variable is not set. "
            "Please set it using: export OPENROUTER_API_KEY='your-key-here'"
        )

    # Warn about Gemini models known limitations
    if any(gemini_model in model_id for gemini_model in GEMINI_MODELS):
        logger.warning(
            f"WARNING: Google Gemini models ({model_id}) may not work reliably with CodeAgent due to "
            "authentication issues when called through OpenRouter. These models require direct Google "
            "authentication which is not compatible with the OpenAI-compatible API endpoint used by CodeAgent. "
            "Consider using alternative models like google/gemma-3-27b-it:free or xiaomi/mimo-v2-flash:free"
        )

    ts_start = dt.datetime.now()
    model_idshort = get_model_idshort(model_id)

    log_event(
        MODEL=model_id,
        EVENT="TASK_START",
        DURATION_SEC=None,
        INPUT_FILE=None,
        OUTPUT_FILE=None,
        STATUS="SUCCESS"
    )

    try:
        # Select model based on provider
        if provider == "ollama":
            # Use LiteLLMModel for Ollama
            model = LiteLLMModel(
                model_id=f"ollama/{model_id}",
                api_base="http://localhost:11434",
                api_key="ollama"  # LiteLLM requires an api_key even for Ollama
            )
        else:
            # Use OpenAIModel for OpenRouter
            model = OpenAIModel(
                model_id=model_id,
                api_base="https://openrouter.ai/api/v1",
                api_key=os.environ["OPENROUTER_API_KEY"],
            )

        result = ""

        with CodeAgent(
            tools=[WebSearchTool()],  # No additional tools - use default CodeAgent tools
                additional_authorized_imports=["sympy", "numpy", "pandas", "matplotlib", "matplotlib.pyplot"],
            model=model
        ) as agent:
            # Use built-in system prompt - no custom prompt override
            result = agent.run(prompt,
                               reset=True,
                               max_steps=max_steps)

        ts_end = dt.datetime.now()
        duration = (ts_end - ts_start).total_seconds()
        log_event(
            MODEL=model_id,
            EVENT="TASK_END",
            DURATION_SEC=f"{duration:.1f}",
            INPUT_FILE=None,
            OUTPUT_FILE=None,
            STATUS="SUCCESS"
        )

        return result

    except KeyError as e:
        if "OPENROUTER_API_KEY" in str(e):
            raise ValueError(
                "OPENROUTER_API_KEY environment variable is not set. "
                "Please set it using: export OPENROUTER_API_KEY='your-key-here'"
            )
        raise
    except Exception as e:
        ts_end = dt.datetime.now()
        duration = (ts_end - ts_start).total_seconds()
        log_event(
            MODEL=model_id,
            EVENT="TASK_END",
            DURATION_SEC=f"{duration:.1f}",
            INPUT_FILE=None,
            OUTPUT_FILE=None,
            STATUS=f"FAILED: {type(e).__name__}"
        )

        # Provide helpful error messages for common issues
        error_msg = str(e)
        if "401" in error_msg and "cookie" in error_msg.lower():
            raise RuntimeError(
                f"Authentication error with model '{model_id}': {error_msg}\n\n"
                "This error often occurs with Google Gemini models which require direct Google authentication. "
                f"The model '{model_id}' may not be compatible with OpenRouter's OpenAI-compatible API endpoint. "
                "Try using: google/gemma-3-27b-it:free, xiaomi/mimo-v2-flash:free, or nex-agi/deepseek-v3.1-nex-n1:free"
            ) from e
        raise RuntimeError(
            f"Error in generating model output:\n{error_msg}"
        ) from e
