tools/export.py

from typing import Dict, List, Any, Optional, Union
import pandas as pd
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
import base64
import io
from pathlib import Path
import json
import datetime

class ExportTools:
    """Tools for exporting data, generating reports, and sending emails."""
    
    def __init__(self, output_directory: str = "./exports"):
        """Initialize with directory for saved exports."""
        self.output_directory = Path(output_directory)
        self.output_directory.mkdir(exist_ok=True, parents=True)
    
    def get_tools(self) -> List[Dict[str, Any]]:
        """Get all available export tools."""
        tools = [
            {
                "name": "generate_report",
                "description": "Generate a report from conversation and results",
                "function": self.generate_report
            },
            {
                "name": "save_results_to_csv",
                "description": "Save query results to a CSV file",
                "function": self.save_results_to_csv
            },
            {
                "name": "send_email",
                "description": "Send results via email",
                "function": self.send_email
            }
        ]
        return tools
    
    def generate_report(self, title: str, content: str, 
                       images: List[str] = None) -> Dict[str, Any]:
        """Generate HTML report from content and images."""
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{title.replace(' ', '_')}_{timestamp}.html"
        file_path = self.output_directory / filename
        
        # Basic HTML template
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>{title}</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 40px; line-height: 1.6; }}
                h1 {{ color: #333366; }}
                .report-container {{ max-width: 800px; margin: 0 auto; }}
                .timestamp {{ color: #666; font-size: 0.8em; }}
                img {{ max-width: 100%; height: auto; margin: 20px 0; }}
            </style>
        </head>
        <body>
            <div class="report-container">
                <h1>{title}</h1>
                <div class="timestamp">Generated on: {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</div>
                <div class="content">
                    {content.replace(chr(10), '<br>')}
                </div>
        """
        
        # Add images if provided
        if images and len(images) > 0:
            html += "<div class='images'>"
            for i, img_base64 in enumerate(images):
                html += f"<img src='data:image/png;base64,{img_base64}' alt='Figure {i+1}'>"
            html += "</div>"
            
        html += """
            </div>
        </body>
        </html>
        """
        
        # Write to file
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(html)
            
        return {
            "success": True,
            "report_path": str(file_path),
            "title": title,
            "timestamp": timestamp
        }
    
    def save_results_to_csv(self, data: List[Dict[str, Any]], 
                           filename: str = None) -> Dict[str, Any]:
        """Save query results to a CSV file."""
        if not data or len(data) == 0:
            return {"success": False, "error": "No data provided"}
            
        # Create DataFrame from data
        df = pd.DataFrame(data)
        
        # Generate filename if not provided
        if not filename:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"query_results_{timestamp}.csv"
            
        # Ensure filename has .csv extension
        if not filename.lower().endswith('.csv'):
            filename += '.csv'
            
        file_path = self.output_directory / filename
        
        # Save to CSV
        df.to_csv(file_path, index=False)
        
        return {
            "success": True,
            "file_path": str(file_path),
            "row_count": len(df),
            "column_count": len(df.columns)
        }
    
    def send_email(self, to_email: str, subject: str, body: str,
                  from_email: str = None, smtp_server: str = None,
                  smtp_port: int = 587, username: str = None,
                  password: str = None, images: List[str] = None) -> Dict[str, Any]:
        """
        Send email with results.
        Note: In production, credentials should be securely managed.
        For demo purposes, this will log the email content instead.
        """
        # For safety in a demo app, don't actually send emails
        # Just log what would be sent and return success
        
        email_content = {
            "to": to_email,
            "subject": subject,
            "body": body[:100] + "..." if len(body) > 100 else body,
            "images": f"{len(images) if images else 0} images would be attached",
            "note": "Email sending is simulated for demo purposes"
        }
        
        # Log the email content
        print(f"SIMULATED EMAIL: {json.dumps(email_content, indent=2)}")
        
        return {
            "success": True,
            "to": to_email,
            "subject": subject,
            "simulated": True,
            "timestamp": datetime.datetime.now().isoformat()
        }