import os
import requests
from rich.console import Console
# from smolagents.agents import ToolCallingAgent
from smolagents import ToolCallingAgent, LiteLLMModel, tool
from smolagents.memory import SystemPromptStep, TaskStep, ActionStep, PlanningStep
from smolagents.monitoring import LogLevel

# For anthropic: change model_id below to 'anthropic/claude-3-5-sonnet-20240620'
model = LiteLLMModel(
    model_id="groq/deepseek-r1-distill-llama-70b",
    api_base="https://api.groq.com/openai/v1",
    api_key=os.environ.get("GROQ_API_KEY", "ERROR"),
    # max_completion_tokens=4096,
)

console = Console()

@tool
def get_joke() -> str:
    """
    Fetches a random joke from the JokeAPI.
    This function sends a GET request to the JokeAPI to retrieve a random joke.
    It handles both single jokes and two-part jokes (setup and delivery).
    If the request fails or the response does not contain a joke, an error message is returned.
    Returns:
        str: The joke as a string, or an error message if the joke could not be fetched.
    """
    url = "https://v2.jokeapi.dev/joke/Any?type=single"

    try:
        response = requests.get(url)
        response.raise_for_status()

        data = response.json()

        if "joke" in data:
            return data["joke"]
        elif "setup" in data and "delivery" in data:
            return f"{data['setup']} - {data['delivery']}"
        else:
            return "Error: Unable to fetch joke."

    except requests.exceptions.RequestException as e:
        return f"Error fetching joke: {str(e)}"


agent = ToolCallingAgent(
    tools=[
        get_joke,
    ],
    model=model,
)


agent.run("Tell me a joke")
agent.replay() # ERROR HERE 

# Soemthing like this might work, but might be missing data (tasks, model_output?)
# work 
console.print("[bold]Summary of agent actions:[/bold]")

for i, step in enumerate(agent.memory.steps):
    class_name = type(step).__name__
    if isinstance(step, SystemPromptStep):
        agent.logger.log_markdown(
            title="System prompt", content=step.system_prompt, level=LogLevel.ERROR
        )
    elif isinstance(step, TaskStep):
        agent.logger.log_task(step.task, "", level=LogLevel.ERROR)
    elif isinstance(step, ActionStep):
        agent.logger.log_rule(f"Step {step.step_number}", level=LogLevel.ERROR)
        agent.logger.log_messages(step.model_input_messages)

        agent.logger.log_markdown(
            title="Agent output:", content=(step.model_output if step.model_output is not None else ""), level=LogLevel.ERROR
        )
    elif isinstance(step, PlanningStep):
        agent.logger.log_rule("Planning step", level=LogLevel.ERROR)
        agent.logger.log_messages(step.model_input_messages, level=LogLevel.ERROR)
        agent.logger.log_markdown(
            title="Agent output:",
            content=step.plan,
            level=LogLevel.ERROR,
        )