import os
import requests
import inspect
import pandas as pd
from dataclasses import asdict
from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, GoogleSearchTool, LiteLLMModel
from  tools import  SafeDuckDuckGoSearchTool, ImageUnderstanding, WikipediaSearch, ExcelReader, CsvReader,ChessSolver, download_files, get_images, FileReader, AudioTransciber, YouTubeVideoUnderstanding, WeatherTool, OCR, TextToImageTool
from pathlib import Path
from dotenv import find_dotenv, load_dotenv
from PIL import Image
import time
from smolagents import stream_to_gradio, Tool

class PowerAgent(CodeAgent):
    def __init__(self, model):
        web_search = DuckDuckGoSearchTool()
        tools = [
            SafeDuckDuckGoSearchTool(),
            ImageUnderstanding(),
            WikipediaSearchTool(),
            FileReader(),
            ExcelReader(),
            CsvReader(),
            AudioTransciber(),
            YouTubeVideoUnderstanding(),
            WeatherTool(),
            ChessSolver(),
            OCR(),
            TextToImageTool()
        ]
        super().__init__(
            tools=tools,
            model=model,
            verbosity_level=3,
            additional_authorized_imports=['pandas', 'numpy', 'io']
        )
        # Add your custom system prompt
        SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. Report your thoughts, and
        finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated
        list of numbers and/or strings.
        If you are asked for a number, don't use comma to write your number neither use units such as $ or
        percent sign unless specified otherwise.
        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the
        digits in plain text unless specified otherwise. Never include currency symbols in the response. 
        If you are asked for a comma separated list, apply the above rules depending of whether the element
        to be put in the list is a number or a string. For question that contain phrases like `what is the number` or 
        `what is the highest number` return just the number, e.g., 2. For questions around currency, 
        include just the number, not the currency sign. Always use `final_answer` tool unless the output is an image,
        then dislay the image.
        """
        self.prompt_templates["system_prompt"] = self.prompt_templates["system_prompt"] + SYSTEM_PROMPT