from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import os
from huggingface_hub import list_models
from Gradio_UI import GradioUI
import baostock as bs
import pandas as pd
import re

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_cutom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

# Weather Data Tool using weatherapi.com API
@tool
def weather_data_tool(location: str, api_key: str) -> str:
    """
    A tool that fetches the current weather data for a given location using the weatherapi.com API.
    
    Args:
        location: The location for which to retrieve weather data (e.g., "London", "Paris", or "New York").
        api_key: d3d95d8e1ada4f48a83112710251202
    
    Returns:
        A string with the current weather details (temperature, condition, etc.) or an error message.
    """
    try:
        # Construct the API URL for current weather data
        url = f"http://api.weatherapi.com/v1/current.json?key={api_key}&q={location}&aqi=no"
        response = requests.get(url)
        if response.status_code != 200:
            return f"Error: Unable to retrieve data. HTTP Status Code: {response.status_code}"          
        data = response.json()       
        # If the response includes an error, return its message
        if "error" in data:
            return f"Error: {data['error'].get('message', 'Unknown error')}"       
        # Extract location and current weather details
        location_info = data.get("location", {})
        current = data.get("current", {})      
        location_name = location_info.get("name", location)
        region = location_info.get("region", "")
        country = location_info.get("country", "")
        temp_c = current.get("temp_c", "N/A")
        condition = current.get("condition", {}).get("text", "N/A")       
        # Build a result string
        result = f"Weather in {location_name}"
        if region:
            result += f", {region}"
        if country:
            result += f", {country}"
        result += f": {temp_c}°C, {condition}."       
        return result
    except Exception as e:
        return f"An exception occurred: {str(e)}"

@tool
def model_download_tool(task: str) -> str:
    """
    This is a tool that returns the most downloaded model of a given task on the Hugging Face Hub.
    It returns the name of the checkpoint.

    Args:
        task: The task for which to get the download count.
    """
    most_downloaded_model = next(iter(list_models(filter=task, sort="downloads", direction=-1)))
    return most_downloaded_model.id

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

@tool
def get_stock_data_baostock(stock_code: str) -> str:
    """Fetch the latest stock data for a given stock code using Baostock.
    Args:
        stock_code: A string representing the stock code (e.g., 'sh.600941' for China Mobile).
    """
    try:
        # Log in to the Baostock system
        lg = bs.login()
        if lg.error_code != '0':
            return f"Failed to log in to Baostock: {lg.error_msg}"

        # Fetch daily K-line data for the stock (defaults to the most recent trading day)
        rs = bs.query_history_k_data_plus(
            code=stock_code,
            fields="date,code,open,high,low,close,volume",  # Fields to retrieve
            frequency="d",  # Daily K-line
            adjustflag="2"  # Adjustment type: 2 for forward adjustment
        )
        if rs.error_code != '0':
            return f"Failed to fetch stock data: {rs.error_msg}"
        # Convert the data to a DataFrame
        data_list = []
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
        df = pd.DataFrame(data_list, columns=rs.fields)
        # Log out of the Baostock system
        bs.logout()
        # Check if the data is empty
        if df.empty:
            return f"No data found for stock code '{stock_code}'."
        # Extract the latest trading day's data
        latest_data = df.iloc[-1]  # Get the last row of data (most recent trading day)
        date = latest_data['date']  # Date
        stock_name = latest_data['code']  # Stock code
        close_price = latest_data['close']  # Closing price
        volume = latest_data['volume']  # Trading volume
        # Construct the result string
        result = f"Stock Code: {stock_name}\nDate: {date}\nClosing Price: {close_price}\nVolume: {volume}"
        return result
    except Exception as e:
        return f"An error occurred while fetching stock data: {str(e)}"
# Example call
print(get_stock_data_baostock('sh.600941'))

final_answer = FinalAnswerTool()

# Import tool from Hub
image_generation_tool = load_tool("m-ric/text-to-image", trust_remote_code=True, token=os.getenv('hf_token'))

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

model = HfApiModel(
    max_tokens=3000,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    #model_id='deepseek-ai/DeepSeek-R1-Distill-Qwen-32B',
    #model_id='mistralai/Mistral-7B-Instruct-v0.3',
    #model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',
    #model_id="google/gemma-2-9b-it",
    custom_role_conversions=None,
    token=os.getenv('hf_token'),
)

agent = CodeAgent(
    model=model,
    tools=[final_answer, image_generation_tool, get_stock_data_baostock, weather_data_tool], ## add your tools here (don't remove final answer)
    max_steps=4,
    verbosity_level=1,
    grammar=None,
    planning_interval=3,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
    add_base_tools=True,
    additional_authorized_imports=["time", "numpy", "pandas","datetime"],
)

GradioUI(agent).launch()