# Importing necessary modules and tools for our application.
# These libraries help us work with code, dates, files, images, and the user interface.
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime      # Helps with dates and times.
import requests      # Lets us make web requests.
import yaml          # Reads YAML files, which are used for configuration.
from tools.final_answer import FinalAnswerTool   # Custom tool to finalize the answer.
from Gradio_UI import GradioUI   # Module for creating a simple web interface.

# -----------------------------
# Tool: comment_code
# -----------------------------
# This function takes a piece of code and adds comments to explain what each part does.
@tool
def comment_code(code: str) -> str:
    """
    A tool that takes a block of code and returns the same code with added comments.
    The comments explain each part of the code in simple terms.
    
    Args:
        code: The input code as a string.
    """
    try:
        # Start with an empty string to collect our commented code.
        commented_code = ""
        # Split the code into lines so we can process each line separately.
        for line in code.split('\n'):
            stripped = line.strip()  # Remove extra spaces from the beginning and end.
            if stripped:
                # Check if the line defines a function.
                if "def " in stripped:
                    commented_code += "# This defines a function\n"
                # Check if the line imports a module.
                elif "import " in stripped:
                    commented_code += "# Importing necessary modules\n"
                # Check if the line is a loop (for or while).
                elif "for " in stripped or "while " in stripped:
                    commented_code += "# Looping through values\n"
                # Check if the line contains a conditional statement.
                elif "if " in stripped:
                    commented_code += "# Checking a condition\n"
                # Check if the line returns a value.
                elif "return " in stripped:
                    commented_code += "# Returning a value from the function\n"
            # Add the original code line after the comment.
            commented_code += line + "\n"
        return commented_code
    except Exception as e:
        # If something goes wrong, return an error message.
        return f"Error processing code: {str(e)}"

# -----------------------------
# Tool: extract_code_from_image
# -----------------------------
# This tool extracts text (code) from an image file, like a screenshot.
@tool
def extract_code_from_image(image_path: str) -> str:
    """
    Extracts text from an image file (a screenshot) that contains code.
    
    Args:
        image_path: The path to the image file.
    """
    try:
        from PIL import Image      # Module for opening and working with images.
        import pytesseract         # OCR tool to extract text from images.
        image = Image.open(image_path)  # Open the image file.
        code_text = pytesseract.image_to_string(image)  # Extract text from the image.
        return code_text
    except Exception as e:
        # Return an error message if extraction fails.
        return f"Error extracting code from image: {str(e)}"

# -----------------------------
# Tool: extract_code_from_file
# -----------------------------
# This tool reads code from a file and returns it as text.
@tool
def extract_code_from_file(file_path: str) -> str:
    """
    Reads code text from a file.
    
    Args:
        file_path: The path to the file containing code.
    """
    try:
        with open(file_path, 'r') as f:   # Open the file in read mode.
            code_text = f.read()          # Read the content of the file.
        return code_text
    except Exception as e:
        # Return an error message if the file cannot be read.
        return f"Error reading code from file: {str(e)}"

# Create an instance of the FinalAnswerTool.
final_answer = FinalAnswerTool()

# -----------------------------
# Set up the language model.
# -----------------------------
# This model will help the agent understand and process code-related tasks.
model = HfApiModel(
    max_tokens=2096,              # Maximum tokens the model can output.
    temperature=0.5,              # Controls randomness in the model's output.
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',   # Specific model to use.
    custom_role_conversions=None, # No custom role conversion needed here.
)

# -----------------------------
# Optionally, load an image generation tool.
# -----------------------------
# This tool is loaded from an external source, though it isn't used directly here.
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# -----------------------------
# Load prompt templates from a configuration file.
# -----------------------------
# This file likely contains text instructions that help guide the agent.
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# -----------------------------
# Set up the CodeAgent.
# -----------------------------
# The agent uses the model and the tools to perform tasks.
agent = CodeAgent(
    model=model,
    tools=[final_answer, comment_code, extract_code_from_image, extract_code_from_file],
    max_steps=6,                    # Limit on how many steps the agent can take.
    verbosity_level=1,              # Level of detail in the logs/output.
    grammar=None,                   # No special grammar rules are applied.
    planning_interval=None,         # No specific planning interval set.
    name="Code Commenter",          # Name of the agent.
    description="An agent that adds explanatory comments to code in simple terms.",
    prompt_templates=prompt_templates  # Prompts to guide the agent's responses.
)

# -----------------------------
# Launch the Gradio user interface.
# -----------------------------
# Gradio creates a simple web interface for users to interact with the agent.
GradioUI(agent).launch()