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

#my imports
import platform
import subprocess
import shlex
from typing import Tuple, Optional
import re

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: #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 ?"

@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_os_info() -> str:
    """A tool that returns info about the OS this agent runs on"""
    os_info = platform.system()
    return f"I run on {os_info}"

@tool
def execute_bash_command(command: str) -> Tuple[str, str]:
    """
    Executes a given Bash command in a subprocess.
    
    This is useful for running system-level commands, shell scripts, 
    or utility programs. The command is run with a timeout for safety.

    Args:
        command: The Bash command string to execute (e.g., 'ls -l', 'grep "pattern" file.txt').

    Returns:
        A tuple containing (stdout, stderr). The output is split into 
        standard output and standard error streams.
    """
    try:
        # 1. Safely split the command string into a list of arguments
        # This is a recommended practice when using subprocess.run 
        # with shell=False (which is safer).
        args = shlex.split(command)
        
        # 2. Execute the command
        # - capture_output=True: Captures stdout and stderr
        # - text=True: Decodes stdout/stderr as text (using default locale encoding)
        # - timeout=10: Kills the command after 10 seconds for safety
        result = subprocess.run(
            args, 
            capture_output=True, 
            text=True, 
            check=False, # Don't raise an exception for non-zero exit codes (we handle it below)
            timeout=10
        )
        
        # 3. Format the output
        stdout = result.stdout.strip()
        stderr = result.stderr.strip()
        
        if result.returncode != 0 and not stderr:
             # If return code is non-zero but stderr is empty,
             # the stdout might contain the error message.
             stderr = f"Command failed with exit code {result.returncode}. Output: {stdout}"
             stdout = ""
        elif result.returncode != 0:
             stderr = f"Command failed with exit code {result.returncode}.\n{stderr}"


        return (stdout, stderr)

    except FileNotFoundError:
        return ("", f"Error: Command or executable not found: '{args[0]}'")
    except subprocess.TimeoutExpired:
        return ("", f"Error: Command timed out after 10 seconds.")
    except Exception as e:
        return ("", f"An unexpected error occurred: {e}")

@tool
def jedi_mind_trick(question: str) -> str:
    # it's import to specify the return type
    """A powerful tool that uses the classic Jedi Mind Trick phrase to dismiss queries about droids.
    
     Args:
        question: The user's full question or query, which should be about droids (or C3PO/R2D2).
    """
    
    # Convert the question to lowercase for case-insensitive matching
    lower_question = question.lower()
    
    # Keywords to trigger the mind trick
    # Uses regex word boundaries (\b) for 'droid' to avoid matching 'android', 
    # but allows exact matching for character names.
    droid_keywords = re.compile(r'\b(droid|droids)\b|c3po|r2d2|bb8')
    
    if droid_keywords.search(lower_question):
        # The trick works!
        return "These aren't the droids you are looking for"
    else:
        # For non-droid questions, return a neutral response
        return "You don't need to use the Force for this query."

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=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, get_current_time_in_timezone, get_os_info, jedi_mind_trick, DuckDuckGoSearchTool()], ## 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()