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

# Below is an example of a tool that does nothing. Amaze us with your creativity!
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

@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:
        tz = pytz.timezone(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 search_semantic_scholar(query: str, 
                            minCitationCount: int = 100, 
                            year: str = None, 
                            fieldsOfStudy: str = None, 
                            openAccessPdf: bool = False) -> dict:
    """Tool: Search Scientific Publications via Semantic Scholar API.
    
    This tool finds scientific publications by querying the Semantic Scholar API.
    It returns metadata about papers matching your search criteria. You can use this tool to
    quickly gather research-based information on your question, and for each used article,
    always reference its Title and URL in the final answer.
    
    **Query Parameter Syntax:**  
    - The `query` parameter accepts a plain-text string containing search keywords.
    - No special query syntax is supported.
    - Hyphenated terms yield no matches; use spaces to separate keywords.
    - Examples: "generative ai", "vitamin C marathon recovery"
    
    **Optional Filters:**  
    - `minCitationCount`: Restricts results to only include papers with at least this many citations. Default is 100.
    - `year`: Restricts results to a specific publication year or a range (e.g., "2016-2020"). Default is any.
    - `fieldsOfStudy`: Restricts results to papers in the specified fields of study as a comma-separated list.  
      Example: "Computer Science,Mathematics,Physics"
    - `openAccessPdf`: If True, restricts results to only include papers that have an open access PDF available.
    
    **Returns:**  
    A JSON dictionary containing the search results from the Semantic Scholar API. Expected structure:
    
    {
        "total": int,             // Total number of matching papers.
        "offset": int,            // Pagination offset used for the current query.
        "data": [                 // Array of paper objects, each with the following properties:
            {
                "paperId": str,   // Unique identifier for the paper (always returned).
                "title": str,     // Title of the paper. Use this as reference in the final answer.
                "abstract": str,  // Abstract summarizing the paper content.
                "year": int,      // Publication year.
                "url": str        // URL to the paper details or full text. Use this as reference in the final answer.
            },
            ...
        ]
    }
    
    In case of an error, the returned dict contains:
    
    {
        "error": int,           // HTTP error status code.
        "message": str          // Description of the error.
    }
    
    **Final Answer Guidance:**  
    For every article used in the final answer, always include its Title and URL as reference.
    
    Args:
        query: Plain-text search query.
        minCitationCount: Minimum number of citations required. Default is 100.
        year: Specific publication year or range (e.g., "2016-2020"). Default is any.
        fieldsOfStudy: Comma-separated list of fields of study.  
                       Example: "Computer Science,Mathematics,Physics". Default is any.
        openAccessPdf: If True, restricts results to papers with an open access PDF available.
      
    Returns:
        dict: JSON response from the Semantic Scholar API.
    """
    base_url = "https://api.semanticscholar.org/graph/v1/paper/search"
    params = {
        "query": query,
        "minCitationCount": str(minCitationCount),
        "fields": "title,abstract,year,url",  # Always return these fields (paperId is always returned)
        "limit": 20
    }
    if year:
        params["year"] = year
    if fieldsOfStudy:
        params["fieldsOfStudy"] = fieldsOfStudy
    if openAccessPdf:
        params["openAccessPdf"] = ""
    
    response = requests.get(base_url, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": response.status_code, "message": response.text}

final_answer = FinalAnswerTool()

# 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=4192,
    temperature=0.5,
    #model_id='Qwen/Qwen2.5-Coder-32B-Instruct', # it is possible that this model may be overloaded
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
    #model_id='deepseek-ai/DeepSeek-R1-Distill-Qwen-32B',
    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, get_current_time_in_timezone, search_semantic_scholar],
    max_steps=10,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()