from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

import requests

import requests
from duckduckgo_search import DDGS

# ----- WeatherTool -----
class WeatherTool:
    def __init__(self, api_key: str):
        self.api_key = api_key

    def __call__(self, city: str):
        try:
            url = f"https://api.openweathermap.org/data/2.5/weather?q={city}&appid={self.api_key}&units=metric"
            response = requests.get(url)
            data = response.json()

            if response.status_code != 200:
                return f"Error: {data.get('message', 'Unable to fetch weather data')}"

            weather = data["weather"][0]["description"]
            temp = data["main"]["temp"]
            feels_like = data["main"]["feels_like"]
            humidity = data["main"]["humidity"]
            wind_speed = data["wind"]["speed"]

            return (f"Weather in {city.title()}:\n"
                    f"- Condition: {weather}\n"
                    f"- Temperature: {temp}°C (Feels like {feels_like}°C)\n"
                    f"- Humidity: {humidity}%\n"
                    f"- Wind Speed: {wind_speed} m/s")
        except Exception as e:
            return f"Exception during weather fetch: {str(e)}"

# ----- VisitWebpageTool -----
class VisitWebpageTool:
    def __call__(self, url: str):
        return f"Pretending to visit webpage: {url}"

# ----- CustomDuckDuckGoSearchTool -----
class CustomDuckDuckGoSearchTool:
    def __call__(self, query: str, max_results: int = 3):
        try:
            with DDGS() as ddgs:
                results = []
                for r in ddgs.text(query):
                    results.append(r)
                    if len(results) >= max_results:
                        break
            return results
        except Exception as e:
            return f"Search error: {str(e)}"

# ----- FinalAnswerTool -----
class FinalAnswerTool:
    def __call__(self, result: str):
        return f"\n Final Answer:\n{result}"

# ----- DummyModel (if needed) -----
class DummyModel:
    def call(self, input_text):
        return f"Model processed: {input_text}"

# ----- ToolCallingAgent -----
class ToolCallingAgent:
    def __init__(self, tools, final_tool, model, max_steps=10):
        self.tools = tools
        self.final_tool = final_tool
        self.model = model
        self.max_steps = max_steps

    def run(self, input_text):
        print(f"\nAgent received: {input_text}")
        for tool in self.tools:
            try:
                result = tool(input_text)
                if result:
                    final_output = self.final_tool(str(result))
                    print(final_output)
                    return final_output
            except Exception as e:
                print(f"\n[{tool.__class__.__name__} Error]: {str(e)}")

# Initialize tools
weather_tool = WeatherTool(api_key="HF_Token")
visit_webpage = VisitWebpageTool()
duckduckgo_tool = CustomDuckDuckGoSearchTool()
final_answer = FinalAnswerTool()
model = DummyModel()

# Initialize and run the agent
web_agent = ToolCallingAgent(
    tools=[duckduckgo_tool, visit_webpage, weather_tool],
    final_tool=final_answer,
    model=model,
    max_steps=10
)

# Example query
web_agent.run("Tokyo")  # Can try a city or any other search term


# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()