main.py

from concurrent.futures import ThreadPoolExecutor
from io import BytesIO, StringIO
from pathlib import Path
from fastapi import FastAPI, Response
import pandas as pd
from aiven_keep_alive_service import ping_aiven_projects
from cronjob_service import init_scheduler
from dotenv import load_dotenv
from keep_projects_alive_service import ping_all_projects
from mongodb_keep_alive_service import ping_mongodb_projects
from neondb_postgres_keep_alive_service import ping_neondb_projects
from pinecone_keep_alive_service import ping_all_pinecone_indexes
from python_code_interpreter_service import execute_python_code
from redis_keep_alive_service import ping_all_redis_projects
from supabase_keep_alive_service import ping_all_supabase_projects
import os
from fastapi import FastAPI, HTTPException
from typing import Any, Dict
import traceback
import warnings
from fastapi.concurrency import run_in_threadpool
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
import logging
import base64

load_dotenv()

VALID_TOKEN = os.getenv("CODE_INTERPRETER_TOKEN")

warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)

# get max workers from system
max_workers = os.cpu_count() - 2
executor = ThreadPoolExecutor(max_workers=max_workers)

app = FastAPI()

# Setup bearer token security
security = HTTPBearer()

scheduler = init_scheduler()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@app.get("/")
def read_root():
    return {"message": "Cronjob Service is running"}

@app.get("/ping")
def health_check():
    return {"status": "ok"}


@app.get("/ping_all_projects")
def ping_all_projects_endpoint():
    result = ping_all_projects()
    return result



async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    Verify the provided bearer token matches our dummy token
    
    Args:
        credentials: HTTPAuthorizationCredentials containing the token
        
    Returns:
        bool: True if token is valid
        
    Raises:
        HTTPException: 401 if token is invalid
    """
    if credentials.credentials != VALID_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return True


@app.post("/python/execute")
async def execute_code(
    code_request: dict,
    token_valid: bool = Depends(verify_token)
) -> Dict[str, Any]:
    """
    Execute Python code and return the results including parsed data from generated files
    
    Args:
        code_request: Dictionary containing 'code' key with the Python code to execute
        
    Returns:
        dict: Dictionary containing execution results with parsed data
    """
    if not token_valid:
        raise HTTPException(status_code=401, detail="Not authenticated")
    
    if "code" not in code_request:
        raise HTTPException(status_code=400, detail="Code field is required")
    
    code = code_request["code"]
    csv_url = code_request.get("csv_url")
    df = pd.read_csv(csv_url) if csv_url else None
    
    logger.info(f"Executing Python code: {code[:100]}...")
    if csv_url:
        logger.info(f"Using CSV from URL: {csv_url}")
    
    try:
        # Run the code execution
        result = await run_in_threadpool(execute_python_code, code, df)
        
        # Prepare response structure
        response_content = {
            "success": result['error'] is None,
            "output": result['output'],
            "plots": [],
            "html_charts": result.get('html_charts', []),
            "generated_data": []
        }
        
        # Process plots if they exist
        if result.get('plots'):
            response_content["plots"] = [
                {
                    "image": plot_data,
                    "filename": f"chart_{index}.png",
                    "index": index
                }
                for index, plot_data in enumerate(result['plots'])
            ]
        
        # Process generated files - DIRECTLY USE PARSED DATA
        if result.get('generated_files'):
            response_content["generated_data"] = []
            for file in result['generated_files']:
                # Create the file data structure
                file_data = {
                    "filename": file['filename'],
                    "filetype": file['filetype'],
                    "path": file.get('path', ''),
                    "size": file.get('size', 0)
                }
                
                # Add parsed data if available
                if 'data' in file:
                    file_data.update({
                        "data": file['data'],
                        "columns": file.get('columns', []),
                        "dtypes": file.get('dtypes', {}),
                        "shape": file.get('shape', [0, 0])
                    })
                else:
                    file_data["error"] = "No parsed data available"
                
                response_content["generated_data"].append(file_data)
        
        # Add error information if exists
        if result.get('error'):
            response_content["error"] = result['error']
        
        # Clean up temporary files
        data_dir = Path("data")
        if data_dir.exists():
            for file in data_dir.glob('*'):
                try:
                    file.unlink()
                except:
                    pass
        
        return response_content
        
    except Exception as e:
        logger.error(f"Error executing code: {str(e)}")
        logger.error(traceback.format_exc())
        raise HTTPException(
            status_code=500, 
            detail={
                "error": str(e),
                "traceback": traceback.format_exc()
            }
        )
    
       
        
@app.post("/html/get-html")
async def get_html(
    html_request: dict,
    token_valid: bool = Depends(verify_token)
) -> Dict[str, Any]:
    """
    Get HTML content for a specific chart
    
    Args:
        html_request: Dictionary containing 'filename' key with the chart filename
        
    Returns:
        dict: Dictionary containing the HTML content of the requested chart
    """
    if not token_valid:
        raise HTTPException(status_code=401, detail="Not authenticated")
    
    if "filename" not in html_request:
        raise HTTPException(status_code=400, detail="Filename field is required")
    
    filename = html_request["filename"]
    print(f"Request Body: {html_request}")
    print(f"Received request for HTML file: {filename}")
    # Check if file exists
    file_path = Path("charts") / filename
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found")
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
            
        # After reading the file, we can remove the file
        os.remove(file_path)
        
        # Validate the HTML content
        if not html_content:
            raise HTTPException(status_code=400, detail="Empty HTML content")
        
        # Return the HTML content
        return Response(content=html_content, media_type="text/html")

    except Exception as e:
        raise HTTPException(
            status_code=500, 
            detail={
                "error": str(e),
                "traceback": traceback.format_exc()
            }
        )