from smolagents import tool
import requests
import json
import datetime
import os
import base64
from typing import List, Optional, Dict, Any
import pandas as pd
import matplotlib.pyplot as plt
import io


@tool
def web_scrape(url: str) -> str:
    """Scrapes the content from a specified URL.
    
    Args:
        url: The URL to scrape content from.
    """
    try:
        response = requests.get(url, headers={
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
        response.raise_for_status()
        return response.text
    except Exception as e:
        return f"Error scraping {url}: {str(e)}"


@tool
def extract_structured_data(text: str, schema: str) -> str:
    """Extracts structured data from text based on a provided schema.
    
    Args:
        text: The text to extract data from.
        schema: JSON schema describing the data structure to extract.
    """
    try:
        # In a real implementation, you might use regex, NLP, or ML models
        # This is a placeholder for demonstrating the concept
        return f"Extracted structured data according to schema: {schema}"
    except Exception as e:
        return f"Error extracting structured data: {str(e)}"


@tool
def data_visualization(data: str, chart_type: str, title: str = "Data Visualization") -> str:
    """Creates a data visualization from structured data.
    
    Args:
        data: JSON string or CSV text with the data to visualize.
        chart_type: Type of chart to create (bar, line, scatter, pie).
        title: Title for the visualization.
    """
    try:
        # Parse the input data
        try:
            # Try parsing as JSON first
            data_parsed = json.loads(data)
            df = pd.DataFrame(data_parsed)
        except:
            # If not JSON, try as CSV
            csv_data = io.StringIO(data)
            df = pd.DataFrame.from_records(pd.read_csv(csv_data))
        
        # Create appropriate visualization
        plt.figure(figsize=(10, 6))
        
        if chart_type.lower() == 'bar':
            df.plot(kind='bar')
        elif chart_type.lower() == 'line':
            df.plot(kind='line')
        elif chart_type.lower() == 'scatter':
            # Assuming first two columns are x and y
            columns = df.columns
            if len(columns) >= 2:
                plt.scatter(df[columns[0]], df[columns[1]])
            else:
                return "Need at least two columns for scatter plot"
        elif chart_type.lower() == 'pie':
            # Assuming first column is labels, second is values
            columns = df.columns
            if len(columns) >= 2:
                plt.pie(df[columns[1]], labels=df[columns[0]], autopct='%1.1f%%')
            else:
                return "Need at least two columns for pie chart"
        else:
            return f"Unsupported chart type: {chart_type}"
        
        plt.title(title)
        
        # Save to bytes buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        
        # Convert to base64 for embedding in HTML or returning
        img_str = base64.b64encode(buf.read()).decode('utf-8')
        
        # Return a reference or small thumbnail
        return f"Visualization created successfully. Image data (base64): {img_str[:30]}..."
    except Exception as e:
        return f"Error creating visualization: {str(e)}"


@tool
def code_refactor(code: str, language: str, optimization: str) -> str:
    """Refactors code based on specified optimization criteria.
    
    Args:
        code: The source code to refactor.
        language: Programming language of the code.
        optimization: Type of optimization to perform (performance, readability, security).
    """
    try:
        # In a real implementation, you'd use language-specific tools or ML models
        # This is a placeholder for demonstrating the concept
        if optimization.lower() == 'performance':
            return f"Code refactored for performance: \n```{language}\n# Performance optimized\n{code}\n```"
        elif optimization.lower() == 'readability':
            return f"Code refactored for readability: \n```{language}\n# Readability optimized\n{code}\n```"
        elif optimization.lower() == 'security':
            return f"Code refactored for security: \n```{language}\n# Security optimized\n{code}\n```"
        else:
            return f"Unsupported optimization type: {optimization}"
    except Exception as e:
        return f"Error refactoring code: {str(e)}"


@tool
def api_interaction(endpoint: str, method: str = "GET", params: Optional[str] = None, headers: Optional[str] = None) -> str:
    """Interacts with an API endpoint.
    
    Args:
        endpoint: The API endpoint URL.
        method: HTTP method (GET, POST, PUT, DELETE).
        params: JSON string of parameters or data to send.
        headers: JSON string of headers to include.
    """
    try:
        # Parse headers and params if provided
        headers_dict = json.loads(headers) if headers else {}
        
        if method.upper() == "GET":
            params_dict = json.loads(params) if params else {}
            response = requests.get(endpoint, params=params_dict, headers=headers_dict)
        elif method.upper() == "POST":
            data_dict = json.loads(params) if params else {}
            response = requests.post(endpoint, json=data_dict, headers=headers_dict)
        elif method.upper() == "PUT":
            data_dict = json.loads(params) if params else {}
            response = requests.put(endpoint, json=data_dict, headers=headers_dict)
        elif method.upper() == "DELETE":
            response = requests.delete(endpoint, headers=headers_dict)
        else:
            return f"Unsupported HTTP method: {method}"
        
        response.raise_for_status()
        
        # Try to return JSON if possible, otherwise return text
        try:
            return json.dumps(response.json(), indent=2)
        except:
            return response.text
    except Exception as e:
        return f"Error interacting with API {endpoint}: {str(e)}"


@tool
def natural_language_query(database_description: str, query: str) -> str:
    """Translates a natural language query to structured data operations.
    
    Args:
        database_description: Description of the database schema.
        query: Natural language query about the data.
    """
    try:
        # In a real implementation, you'd use NLP to SQL or similar technology
        # This is a placeholder for demonstrating the concept
        return f"Query translated and executed. Results for: {query}"
    except Exception as e:
        return f"Error processing natural language query: {str(e)}"


@tool
def file_operations(operation: str, file_path: str, content: Optional[str] = None) -> str:
    """Performs operations on files.
    
    Args:
        operation: The operation to perform (read, write, append, list).
        file_path: Path to the file or directory.
        content: Content to write or append (only for write/append operations).
    """
    try:
        if operation.lower() == 'read':
            with open(file_path, 'r') as file:
                return file.read()
        elif operation.lower() == 'write':
            if content is None:
                return "Content must be provided for write operation"
            with open(file_path, 'w') as file:
                file.write(content)
            return f"Content written to {file_path}"
        elif operation.lower() == 'append':
            if content is None:
                return "Content must be provided for append operation"
            with open(file_path, 'a') as file:
                file.write(content)
            return f"Content appended to {file_path}"
        elif operation.lower() == 'list':
            if os.path.isdir(file_path):
                return str(os.listdir(file_path))
            else:
                return f"{file_path} is not a directory"
        else:
            return f"Unsupported file operation: {operation}"
    except Exception as e:
        return f"Error performing file operation: {str(e)}"


@tool
def semantic_search(corpus: str, query: str, top_k: int = 3) -> str:
    """Performs semantic search on a corpus of text.
    
    Args:
        corpus: The text corpus to search within (could be a large text or list of documents).
        query: The search query.
        top_k: Number of top results to return.
    """
    try:
        # In a real implementation, you'd use embedding models and vector similarity
        # This is a placeholder for demonstrating the concept
        results = [
            {"text": f"Result {i} for query: {query}", "score": (top_k - i) / top_k}
            for i in range(1, top_k + 1)
        ]
        return json.dumps(results, indent=2)
    except Exception as e:
        return f"Error performing semantic search: {str(e)}"


@tool
def weather_forecast(location: str) -> str:
    """Fetches weather forecast for a specified location.
    
    Args:
        location: The location to get weather forecast for (city name or coordinates).
    """
    try:
        # In a real implementation, you'd connect to a weather API
        # This is a placeholder for demonstrating the concept
        return f"Weather forecast for {location}: Sunny with a chance of AI"
    except Exception as e:
        return f"Error fetching weather forecast: {str(e)}"


@tool 
def task_scheduler(task: str, schedule_time: str, priority: int = 1) -> str:
    """Schedules a task to be performed at a specified time.
    
    Args:
        task: Description of the task to be scheduled.
        schedule_time: Time to schedule the task (ISO format).
        priority: Priority level of the task (1-5, where 1 is highest).
    """
    try:
        # Parse the schedule time
        schedule_datetime = datetime.datetime.fromisoformat(schedule_time)
        
        # In a real implementation, you'd connect to a scheduling system
        # This is a placeholder for demonstrating the concept
        return f"Task '{task}' scheduled for {schedule_datetime} with priority {priority}"
    except Exception as e:
        return f"Error scheduling task: {str(e)}"