report_service.py

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import datetime as dt
import os
import uuid
import json
import logging
import sys
import traceback
import shutil
import ast  # <--- Added for syntax checking
from io import StringIO
from typing import List, Dict, Any, Optional
from pydantic import BaseModel
from starlette.concurrency import run_in_threadpool
import numpy as np

# --- Internal Imports ---
from supabase_service import upload_file_to_supabase

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("Report_Generator")

# Configure Matplotlib backend for thread safety (Headless)
os.environ['MPLBACKEND'] = 'agg'

# Disable plt.show globally to prevent blocking
plt.show = lambda: None

# ==============================================================================
# REQUEST MODELS
# ==============================================================================
class ReportRequest(BaseModel):
    code: str
    csv_url: str
    chat_id: str

# ==============================================================================
# RESPONSE MODELS
# ==============================================================================

class FileProps(BaseModel):
    fileName: str
    filePath: str
    fileType: str  # 'csv' | 'image'

class Files(BaseModel):
    csv_files: List[FileProps]
    image_files: List[FileProps]

class FileBoxProps(BaseModel):
    files: Files

# ==============================================================================
# EXECUTION ENGINE
# ==============================================================================

class PythonREPL:
    """Secure Python REPL with file generation tracking"""
    
    def __init__(self, df: pd.DataFrame):
        self.df = df
        # Create a unique directory for this execution to avoid thread collisions
        self.output_dir = os.path.abspath(f'generated_outputs/{uuid.uuid4()}')
        os.makedirs(self.output_dir, exist_ok=True)
        
        self.local_env = {
            "pd": pd,
            "df": self.df.copy(),
            "plt": plt,
            "os": os,
            "uuid": uuid,
            "sns": sns,
            "json": json,
            "dt": dt,
            "np": np,
            "output_dir": self.output_dir # INJECT output_dir so code can use it
        }
        
    def execute(self, code: str) -> Dict[str, Any]:
        logger.info(f'Executing code in: {self.output_dir}')
        old_stdout = sys.stdout
        sys.stdout = mystdout = StringIO()
        
        file_tracker = {
            'csv_files': set(),
            'image_files': set()
        }

        # Wrap code to ensure non-interactive plotting inside the exec scope
        wrapped_code = f"""
import matplotlib.pyplot as plt
plt.switch_backend('agg')
{code}
plt.close('all')
"""
        error = False
        error_msg = None

        try:
            # --- IMPROVEMENT: Syntax Validation ---
            # This checks if the generated code is valid Python before running it.
            # It catches the "missing parenthesis" error gracefully.
            try:
                ast.parse(wrapped_code)
            except SyntaxError as e:
                raise SyntaxError(f"Generated code is incomplete or invalid: {e}")

            # Execute the code
            exec(wrapped_code, self.local_env)
            
            # Check for generated files in the specific directory
            if os.path.exists(self.output_dir):
                for fname in os.listdir(self.output_dir):
                    if fname.endswith('.csv'):
                        file_tracker['csv_files'].add(fname)
                    elif fname.lower().endswith(('.png', '.jpg', '.jpeg')):
                        file_tracker['image_files'].add(fname)
            
        except Exception as e:
            error = True
            error_msg = traceback.format_exc()
            
            # --- IMPROVEMENT: Better Logging ---
            # Log the code causing the error so you can debug the LLM prompt
            logger.error("============= CODE EXECUTION FAILED =============")
            logger.error(f"Error Message: {str(e)}")
            logger.error("---------------- FAILED CODE ----------------")
            for i, line in enumerate(wrapped_code.split('\n')):
                logger.error(f"{i+1}: {line}")
            logger.error("=================================================")
            
        finally:
            sys.stdout = old_stdout
            
        return {
            "output": mystdout.getvalue(),
            "error": error,
            "error_message": error_msg,
            "output_dir": self.output_dir,
            "files": {
                "csv": [os.path.join(self.output_dir, f) for f in file_tracker['csv_files']],
                "images": [os.path.join(self.output_dir, f) for f in file_tracker['image_files']]
            }
        }
    
    def cleanup(self):
        """Remove the temporary directory"""
        if os.path.exists(self.output_dir):
            try:
                shutil.rmtree(self.output_dir)
            except Exception as e:
                logger.warning(f"Cleanup failed: {e}")

# ==============================================================================
# MAIN LOGIC (Synchronous Worker)
# ==============================================================================

def _generate_report_sync(code: str, csv_url: str, chat_id: str) -> FileBoxProps:
    """
    Blocking worker function.
    1. Reads CSV.
    2. Runs Code.
    3. Uploads files synchronously.
    """
    repl = None
    csv_props = []
    image_props = []
    
    try:
        # 1. Load Data
        df = pd.read_csv(csv_url)
        
        # 2. Initialize REPL
        repl = PythonREPL(df)
        
        # 3. Execute Code
        result = repl.execute(code)
        
        if result['error']:
            # Log is already handled in repl.execute
            return FileBoxProps(files=Files(csv_files=[], image_files=[]))

        # 4. Process & Upload CSVs
        for csv_path in result['files']['csv']:
            if os.path.exists(csv_path):
                file_name = os.path.basename(csv_path)
                unique_name = f"{uuid.uuid4()}_{file_name}"
                try:
                    # Sync call - NO AWAIT
                    public_url = upload_file_to_supabase(
                        file_path=csv_path, 
                        file_name=unique_name, 
                        chat_id=chat_id
                    )
                    csv_props.append(FileProps(
                        fileName=file_name,
                        filePath=public_url,
                        fileType="csv"
                    ))
                except Exception as e:
                    logger.error(f"Failed upload CSV {file_name}: {e}")

        # 5. Process & Upload Images
        for img_path in result['files']['images']:
            if os.path.exists(img_path):
                file_name = os.path.basename(img_path)
                unique_name = f"{uuid.uuid4()}_{file_name}"
                try:
                    # Sync call - NO AWAIT
                    public_url = upload_file_to_supabase(
                        file_path=img_path, 
                        file_name=unique_name, 
                        chat_id=chat_id
                    )
                    image_props.append(FileProps(
                        fileName=file_name,
                        filePath=public_url,
                        fileType="image"
                    ))
                except Exception as e:
                    logger.error(f"Failed upload Image {file_name}: {e}")

        return FileBoxProps(
            files=Files(
                csv_files=csv_props,
                image_files=image_props
            )
        )

    except Exception as e:
        logger.error(f"System Error in Report Generator: {e}")
        return FileBoxProps(files=Files(csv_files=[], image_files=[]))
        
    finally:
        # 6. Cleanup local files
        if repl:
            repl.cleanup()

# ==============================================================================
# ASYNC WRAPPER
# ==============================================================================

async def execute_report_generation(code: str, csv_url: str, chat_id: str) -> FileBoxProps:
    """
    Async entry point that runs the blocking logic in a separate thread.
    This enables concurrency.
    """
    return await run_in_threadpool(
        _generate_report_sync,
        code=code,
        csv_url=csv_url,
        chat_id=chat_id
    )