src/tools/enhanced_tools.py

import os
from typing import List, Dict, Any, Optional
import tempfile
import re
import json
import requests
from urllib.parse import urlparse
try:
    import pytesseract
    PYTESSERACT_AVAILABLE = True
except ImportError:
    PYTESSERACT_AVAILABLE = False

from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
import cmath
import pandas as pd
import uuid
import numpy as np
from langchain_core.tools import tool

from src.tools.image_processing import decode_image, encode_image, save_image

### =============== MATHEMATICAL TOOLS =============== ###

@tool
def multiply(a: float, b: float) -> float:
    """Multiplies two numbers (a * b).

    Use this tool for basic multiplication.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        The product of the two numbers.
    """
    return a * b

@tool
def add(a: float, b: float) -> float:
    """Adds two numbers (a + b).

    Use this tool for basic addition.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        The sum of the two numbers.
    """
    return a + b

@tool
def subtract(a: float, b: float) -> float:
    """Subtracts two numbers (a - b).

    Use this tool for basic subtraction.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        The difference between the two numbers.
    """
    return a - b

@tool
def divide(a: float, b: float) -> float:
    """Divides two numbers (a / b).

    Use this tool for basic division. Handles division by zero by raising a ValueError.

    Args:
        a (float): The dividend.
        b (float): The divisor. Must not be zero.

    Returns:
        The quotient of the two numbers.

    Raises:
        ValueError: If the divisor `b` is zero.
    """
    if b == 0:
        raise ValueError("Cannot divided by zero.")
    return a / b

@tool
def modulus(a: int, b: int) -> int:
    """Calculates the remainder of a division (a % b).

    Use this tool to find the modulus of two integers.

    Args:
        a (int): The dividend.
        b (int): The divisor.

    Returns:
        The remainder of the division.
    """
    return a % b

@tool
def power(a: float, b: float) -> float:
    """Raises a number to a specified power (a^b).

    Use this tool for exponentiation.

    Args:
        a (float): The base number.
        b (float): The exponent.

    Returns:
        The result of `a` raised to the power of `b`.
    """
    return a**b

@tool
def square_root(a: float) -> float | complex:
    """Calculates the square root of a number.

    Use this tool to find the square root of a given number. It can handle both positive and negative numbers.

    Args:
        a (float): The number to find the square root of.

    Returns:
        The square root of the number. Returns a complex number if the input is negative.
    """
    if a >= 0:
        return a**0.5
    return cmath.sqrt(a)

### =============== DOCUMENT PROCESSING TOOLS =============== ###

@tool
def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
    """Saves the provided content to a file and returns the file path.

    This tool is useful for persisting information that needs to be accessed later,
    or for creating files that can then be processed by other tools.
    If no filename is provided, a unique temporary file will be created.

    Args:
        content (str): The string content to be saved to the file.
        filename (str, optional): The desired name of the file. If not provided,
                                  a random name will be generated.

    Returns:
        A string indicating the path where the file was saved.
    """
    temp_dir = tempfile.gettempdir()
    if filename is None:
        temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
        filepath = temp_file.name
    else:
        filepath = os.path.join(temp_dir, filename)

    with open(filepath, "w") as f:
        f.write(content)

    return f"File saved to {filepath}. You can read this file to process its contents."

@tool
def read_file(file_path: str) -> str:
    """Reads the contents of a file at the specified path.

    This tool is useful for reading files that have been downloaded, attached, or previously saved.
    It handles text files and returns their contents as a string.

    Args:
        file_path (str): The absolute path to the file to read.

    Returns:
        The contents of the file as a string.
        Returns an error message if the file cannot be read.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return f"File contents:\n{content}"
    except FileNotFoundError:
        return f"Error: File not found at path: {file_path}"
    except Exception as e:
        return f"Error reading file: {str(e)}"

@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """Downloads a file from a given URL and saves it to a temporary location.

    This tool is essential for retrieving external files, such as datasets, images,
    or documents, from the internet for further processing.
    If no filename is provided, the tool will attempt to infer it from the URL
    or generate a unique name.

    Args:
        url (str): The URL of the file to download.
        filename (str, optional): The desired name for the downloaded file.
                                  If not provided, a name will be generated.

    Returns:
        A string indicating the path where the file was downloaded.
        Returns an error message if the download fails.
    """
    try:
        # Parse URL to get filename if not provided
        if not filename:
            path = urlparse(url).path
            filename = os.path.basename(path)
            if not filename:
                filename = f"downloaded_{uuid.uuid4().hex[:8]}"

        # Create temporary file
        temp_dir = tempfile.gettempdir()
        filepath = os.path.join(temp_dir, filename)

        # Download the file
        response = requests.get(url, stream=True)
        response.raise_for_status()

        # Save the file
        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        return f"File downloaded to {filepath}. You can read this file to process its contents."
    except Exception as e:
        return f"Error downloading file: {str(e)}"

@tool
def extract_text_from_image(image_path: str) -> str:
    """Extracts text from an image using Optical Character Recognition (OCR).

    This tool is useful for converting images containing text (e.g., scanned documents,
    screenshots) into machine-readable text. It relies on the `pytesseract` library.

    Args:
        image_path (str): The absolute path to the image file.

    Returns:
        A string containing the extracted text from the image.
        Returns an error message if `pytesseract` is not available or if extraction fails.
    """
    try:
        # Open the image
        image = Image.open(image_path)

        # Extract text from the image
        if PYTESSERACT_AVAILABLE:
            text = pytesseract.image_to_string(image)
            return f"Extracted text from image:\n\n{text}"
        else:
            return "Error: pytesseract not available. Cannot extract text from image."
    except Exception as e:
        return f"Error extracting text from image: {str(e)}"

@tool
def analyze_csv_file(file_path: str, query: str) -> str:
    """Analyzes a CSV file and provides insights based on a given query.

    This tool uses the pandas library to read and perform basic analysis on CSV files.
    It can provide summary statistics, column information, and other data-driven insights.

    Args:
        file_path (str): The absolute path to the CSV file.
        query (str): A natural language question about the data in the CSV file.
                     (Note: The `query` parameter is currently for context and does not
                     directly influence the analysis performed by this tool, which provides
                     general summary statistics. For more specific queries, use `execute_code_multilang`
                     with Python code to analyze the DataFrame.)

    Returns:
        A string containing a summary of the CSV file, including its dimensions,
        column names, and descriptive statistics.
        Returns an error message if the file cannot be read or analyzed.
    """
    try:
        # Read the CSV file
        df = pd.read_csv(file_path)

        # Run various analyses based on the query
        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"

        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())

        return result

    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"

@tool
def analyze_excel_file(file_path: str, query: str) -> str:
    """Analyzes an Excel file and provides insights based on a given query.

    Similar to `analyze_csv_file`, this tool uses the pandas library to read and
    perform basic analysis on Excel files (e.g., .xlsx, .xls).
    It can provide summary statistics, column information, and other data-driven insights.

    Args:
        file_path (str): The absolute path to the Excel file.
        query (str): A natural language question about the data in the Excel file.
                     (Note: The `query` parameter is currently for context and does not
                     directly influence the analysis performed by this tool, which provides
                     general summary statistics. For more specific queries, use `execute_code_multilang`
                     with Python code to analyze the DataFrame.)

    Returns:
        A string containing a summary of the Excel file, including its dimensions,
        column names, and descriptive statistics.
        Returns an error message if the file cannot be read or analyzed.
    """
    try:
        # Read the Excel file
        df = pd.read_excel(file_path)

        # Run various analyses based on the query
        result = (
            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        )
        result += f"Columns: {', '.join(df.columns)}\n\n"

        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())

        return result

    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"

### ============== IMAGE PROCESSING AND GENERATION TOOLS =============== ###

@tool
def analyze_image(image_base64: str) -> Dict[str, Any]:
    """Analyzes basic properties of an image from its base64 representation.

    This tool provides fundamental information about an image, such as its dimensions,
    color mode (e.g., RGB, L), and a basic color analysis (average colors, dominant color).
    It also generates a small thumbnail preview.

    Args:
        image_base64 (str): The base64 encoded string of the image to be analyzed.

    Returns:
        A dictionary containing:
        - `dimensions`: A tuple (width, height) of the image.
        - `mode`: The color mode of the image.
        - `color_analysis`: A dictionary with average RGB, brightness, and dominant color (for RGB/RGBA images).
        - `thumbnail`: A base64 encoded string of a small thumbnail of the image.
        Returns an error message if the image cannot be decoded or analyzed.
    """
    try:
        img = decode_image(image_base64)
        width, height = img.size
        mode = img.mode

        if mode in ("RGB", "RGBA"):
            arr = np.array(img)
            avg_colors = arr.mean(axis=(0, 1))
            dominant = ["Red", "Green", "Blue"][np.argmax(avg_colors[:3])]
            brightness = avg_colors.mean()
            color_analysis = {
                "average_rgb": avg_colors.tolist(),
                "brightness": brightness,
                "dominant_color": dominant,
            }
        else:
            color_analysis = {"note": f"No color analysis for mode {mode}"}

        thumbnail = img.copy()
        thumbnail.thumbnail((100, 100))
        thumb_path = save_image(thumbnail, "thumbnails")
        thumbnail_base64 = encode_image(thumb_path)

        return {
            "dimensions": (width, height),
            "mode": mode,
            "color_analysis": color_analysis,
            "thumbnail": thumbnail_base64,
        }
    except Exception as e:
        return {"error": str(e)}

@tool
def transform_image(
    image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """Applies various transformations to an image from its base64 representation.

    This versatile tool can perform common image manipulations such as resizing,
    rotating, cropping, flipping, adjusting brightness/contrast, blurring, sharpening,
    and converting to grayscale. It's useful for preparing images for display,
    analysis, or further processing.

    Args:
        image_base64 (str): The base64 encoded string of the input image.
        operation (str): The type of transformation to apply. Supported operations:
                         "resize", "rotate", "crop", "flip", "adjust_brightness",
                         "adjust_contrast", "blur", "sharpen", "grayscale".
        params (Dict[str, Any], optional): A dictionary of parameters specific to the
                                          chosen operation (e.g., `{"width": 200, "height": 150}` for resize,
                                          `{"angle": 45}` for rotate, `{"factor": 1.2}` for brightness).

    Returns:
        A dictionary containing the base64 encoded string of the transformed image.
        Returns an error message if the operation is unknown or if transformation fails.
    """
    try:
        img = decode_image(image_base64)
        params = params or {}

        if operation == "resize":
            img = img.resize(
                (
                    params.get("width", img.width // 2),
                    params.get("height", img.height // 2),
                )
            )
        elif operation == "rotate":
            img = img.rotate(params.get("angle", 90), expand=True)
        elif operation == "crop":
            img = img.crop(
                (
                    params.get("left", 0),
                    params.get("top", 0),
                    params.get("right", img.width),
                    params.get("bottom", img.height),
                )
            )
        elif operation == "flip":
            if params.get("direction", "horizontal") == "horizontal":
                img = img.transpose(Image.FLIP_LEFT_RIGHT)
            else:
                img = img.transpose(Image.FLIP_TOP_BOTTOM)
        elif operation == "adjust_brightness":
            img = ImageEnhance.Brightness(img).enhance(params.get("factor", 1.5))
        elif operation == "adjust_contrast":
            img = ImageEnhance.Contrast(img).enhance(params.get("factor", 1.5))
        elif operation == "blur":
            img = img.filter(ImageFilter.GaussianBlur(params.get("radius", 2)))
        elif operation == "sharpen":
            img = img.filter(ImageFilter.SHARPEN)
        elif operation == "grayscale":
            img = img.convert("L")
        else:
            return {"error": f"Unknown operation: {operation}"}

        result_path = save_image(img)
        result_base64 = encode_image(result_path)
        return {"transformed_image": result_base64}

    except Exception as e:
        return {"error": str(e)}

@tool
def draw_on_image(
    image_base64: str, drawing_type: str, params: Dict[str, Any]
) -> Dict[str, Any]:
    """Draws shapes (rectangle, circle, line) or text onto an image.

    This tool allows you to annotate or modify images by adding geometric shapes or text.
    It's useful for highlighting areas, adding labels, or creating simple graphics.

    Args:
        image_base64 (str): The base64 encoded string of the input image.
        drawing_type (str): The type of drawing to perform. Supported types:
                            "rectangle", "circle", "line", "text".
        params (Dict[str, Any]): A dictionary of parameters specific to the drawing type.
                                 - For "rectangle": `{"left", "top", "right", "bottom", "color", "width"}`
                                 - For "circle": `{"x", "y", "radius", "color", "width"}`
                                 - For "line": `{"start_x", "start_y", "end_x", "end_y", "color", "width"}`
                                 - For "text": `{"x", "y", "text", "color", "font_size"}`

    Returns:
        A dictionary containing the base64 encoded string of the image with the drawing.
        Returns an error message if the drawing type is unknown or if drawing fails.
    """
    try:
        img = decode_image(image_base64)
        draw = ImageDraw.Draw(img)
        color = params.get("color", "red")

        if drawing_type == "rectangle":
            draw.rectangle(
                [params["left"], params["top"], params["right"], params["bottom"]],
                outline=color,
                width=params.get("width", 2),
            )
        elif drawing_type == "circle":
            x, y, r = params["x"], params["y"], params["radius"]
            draw.ellipse(
                (x - r, y - r, x + r, y + r),
                outline=color,
                width=params.get("width", 2),
            )
        elif drawing_type == "line":
            draw.line(
                (
                    params["start_x"],
                    params["start_y"],
                    params["end_x"],
                    params["end_y"],
                ),
                fill=color,
                width=params.get("width", 2),
            )
        elif drawing_type == "text":
            font_size = params.get("font_size", 20)
            try:
                font = ImageFont.truetype("arial.ttf", font_size)
            except IOError:
                font = ImageFont.load_default()
            draw.text(
                (params["x"], params["y"]),
                params.get("text", "Text"),
                fill=color,
                font=font,
            )
        else:
            return {"error": f"Unknown drawing type: {drawing_type}"}

        result_path = save_image(img)
        result_base64 = encode_image(result_path)
        return {"result_image": result_base64}

    except Exception as e:
        return {"error": str(e)}

@tool
def generate_simple_image(
    image_type: str,
    width: int = 500,
    height: int = 500,
    params: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Generates a simple image of a specified type.

    This tool can create basic images like gradients or noise patterns.
    It's useful for generating placeholder images, backgrounds, or visual elements
    for testing or simple graphical needs.

    Args:
        image_type (str): The type of image to generate. Supported types:
                          "gradient" (creates a color gradient),
                          "noise" (creates a random noise pattern).
        width (int): The desired width of the generated image in pixels.
        height (int): The desired height of the generated image in pixels.
        params (Dict[str, Any], optional): Parameters specific to the image type.
                                          - For "gradient": `{"direction": "horizontal"|"vertical",
                                            "start_color": (R,G,B), "end_color": (R,G,B)}`

    Returns:
        A dictionary containing the base64 encoded string of the generated image.
        Returns an error message if the image type is unsupported or generation fails.
    """
    try:
        params = params or {}

        if image_type == "gradient":
            direction = params.get("direction", "horizontal")
            start_color = params.get("start_color", (255, 0, 0))
            end_color = params.get("end_color", (0, 0, 255))

            img = Image.new("RGB", (width, height))
            draw = ImageDraw.Draw(img)

            if direction == "horizontal":
                for x in range(width):
                    r = int(
                        start_color[0] + (end_color[0] - start_color[0]) * x / width
                    )
                    g = int(
                        start_color[1] + (end_color[1] - start_color[1]) * x / width
                    )
                    b = int(
                        start_color[2] + (end_color[2] - start_color[2]) * x / width
                    )
                    draw.line([(x, 0), (x, height)], fill=(r, g, b))
            else:
                for y in range(height):
                    r = int(
                        start_color[0] + (end_color[0] - start_color[0]) * y / height
                    )
                    g = int(
                        start_color[1] + (end_color[1] - start_color[1]) * y / height
                    )
                    b = int(
                        start_color[2] + (end_color[2] - start_color[2]) * y / height
                    )
                    draw.line([(0, y), (width, y)], fill=(r, g, b))

        elif image_type == "noise":
            noise_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
            img = Image.fromarray(noise_array, "RGB")

        else:
            return {"error": f"Unsupported image_type {image_type}"}

        result_path = save_image(img)
        result_base64 = encode_image(result_path)
        return {"generated_image": result_base64}

    except Exception as e:
        return {"error": str(e)}

@tool
def combine_images(
    images_base64: List[str], operation: str, params: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """Combines multiple images into a single image.

    This tool supports various combination operations, such as stacking images
    horizontally or vertically. It's useful for creating collages, visual comparisons,
    or composite images.

    Args:
        images_base64 (List[str]): A list of base64 encoded strings of the input images.
        operation (str): The type of combination to perform. Supported operations:
                         "stack" (stacks images side-by-side or top-to-bottom).
        params (Dict[str, Any], optional): Parameters specific to the operation.
                                          - For "stack": `{"direction": "horizontal"|"vertical"}`

    Returns:
        A dictionary containing the base64 encoded string of the combined image.
        Returns an error message if the operation is unsupported or combination fails.
    """
    try:
        images = [decode_image(b64) for b64 in images_base64]
        params = params or {}

        if operation == "stack":
            direction = params.get("direction", "horizontal")
            if direction == "horizontal":
                total_width = sum(img.width for img in images)
                max_height = max(img.height for img in images)
                new_img = Image.new("RGB", (total_width, max_height))
                x = 0
                for img in images:
                    new_img.paste(img, (x, 0))
                    x += img.width
            else:
                max_width = max(img.width for img in images)
                total_height = sum(img.height for img in images)
                new_img = Image.new("RGB", (max_width, total_height))
                y = 0
                for img in images:
                    new_img.paste(img, (0, y))
                    y += img.height
        else:
            return {"error": f"Unsupported combination operation {operation}"}

        result_path = save_image(new_img)
        result_base64 = encode_image(result_path)
        return {"combined_image": result_base64}

    except Exception as e:
        return {"error": str(e)}