pdf_report_generation_helper.py

"""
PDF Report Generation Helper - Production Ready
Handles PDF generation, upload to Supabase, and cleanup
"""
import os
import uuid
import logging
from typing import List, Dict, Any
from starlette.concurrency import run_in_threadpool

# The Service file defines the Generator class
from pdf_report_generation_service import Q4ReportGenerator

logger = logging.getLogger(__name__)

def _upload_to_supabase_sync(file_path: str, file_name: str, chat_id: str) -> str:
    """
    Synchronous wrapper for Supabase upload.
    Import inside function to avoid potential circular dependencies if supabase_service changes.
    """
    from supabase_service import upload_file_to_supabase
    return upload_file_to_supabase(
        file_path=file_path,
        file_name=file_name,
        chat_id=chat_id
    )

def _generate_pdf_sync(
    generator: Q4ReportGenerator,
    config_content: List[Dict[str, Any]],
    local_file_path: str,
    output_dir: str
) -> bool:
    """Synchronous PDF generation wrapper"""
    return generator.generate(
        sections=config_content,
        filename=os.path.basename(local_file_path),
        output_dir=output_dir,
        cover_page=True
    )

def _cleanup_files(file_paths: List[str]):
    """Clean up temporary files"""
    for file_path in file_paths:
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
                logger.info(f"Cleaned up: {file_path}")
            except Exception as e:
                logger.warning(f"Failed to delete {file_path}: {e}")

def _cleanup_downloaded_images(generator: Q4ReportGenerator, output_dir: str):
    """Clean up images downloaded during PDF generation"""
    if not hasattr(generator, 'downloaded_images'):
        return
    
    for img_uri in generator.downloaded_images:
        try:
            # Convert file:// URI to path
            if img_uri.startswith('file://'):
                # Handle both Unix and Windows paths
                if os.name == 'nt' and img_uri.startswith('file:///'):
                    img_path = img_uri[8:]  # Windows: file:///C:/...
                else:
                    img_path = img_uri[7:]  # Unix: file:///path/...
                
                if os.path.exists(img_path):
                    os.remove(img_path)
                    logger.info(f"Deleted temp image: {img_path}")
        except Exception as e:
            logger.warning(f"Failed to delete image {img_uri}: {e}")
    
    # Try to remove the images directory if empty
    images_dir = os.path.join(output_dir, 'images')
    if os.path.exists(images_dir):
        try:
            if not os.listdir(images_dir):  # Check if empty
                os.rmdir(images_dir)
                logger.info(f"Removed empty images directory: {images_dir}")
        except Exception as e:
            logger.debug(f"Could not remove images directory: {e}")

async def generate_and_upload_report(
    config_content: List[Dict[str, Any]],
    file_name: str,
    chat_id: str,
    title: str = "Analytics Report",
    subtitle: str = "Generated Report",
    author: str = "AI Assistant",
    department: str = "Data Analytics",
    output_dir: str = "temp_reports",
    confidential: bool = True
) -> str:
    """
    Generates a PDF report, uploads it to Supabase, and returns the public URL.
    """
    
    local_file_path = None
    
    try:
        # 1. Setup paths
        os.makedirs(output_dir, exist_ok=True)
        
        if not file_name.endswith('.pdf'):
            file_name += '.pdf'
        
        unique_id = str(uuid.uuid4())[:8]
        local_filename = f"{os.path.splitext(file_name)[0]}_{unique_id}.pdf"
        local_file_path = os.path.join(output_dir, local_filename)
        
        logger.info(f"Starting PDF generation: {local_filename}")
        
        # 2. Initialize Generator
        generator = Q4ReportGenerator(
            title=title,
            subtitle=subtitle,
            author=author,
            department=department,
            confidential=confidential
        )
        
        # 3. Generate PDF (Run in thread pool to avoid blocking)
        success = await run_in_threadpool(
            _generate_pdf_sync,
            generator=generator,
            config_content=config_content,
            local_file_path=local_file_path,
            output_dir=output_dir
        )
        
        if not success:
            raise Exception("PDF generation failed - generator returned False")
        
        if not os.path.exists(local_file_path):
            raise Exception(f"PDF file not found after generation: {local_file_path}")
        
        # 4. Upload to Supabase (Run in thread pool)
        cloud_filename = f"{os.path.splitext(file_name)[0]}_{unique_id}.pdf"
        
        logger.info(f"Uploading {local_filename} to Supabase as {cloud_filename}")
        
        public_url = await run_in_threadpool(
            _upload_to_supabase_sync,
            file_path=local_file_path,
            file_name=cloud_filename,
            chat_id=chat_id
        )
        
        if not public_url:
            raise Exception("Upload failed - no URL returned")
        
        logger.info(f"Upload successful: {public_url}")
        
        return public_url
        
    except Exception as e:
        logger.error(f"Report generation/upload failed: {e}", exc_info=True)
        raise Exception(f"Failed to generate and upload report: {str(e)}")
        
    finally:
        # 5. Cleanup
        try:
            if local_file_path:
                _cleanup_files([local_file_path])
            
            if 'generator' in locals():
                await run_in_threadpool(
                    _cleanup_downloaded_images,
                    generator=generator,
                    output_dir=output_dir
                )
        except Exception as cleanup_error:
            logger.warning(f"Cleanup error (non-fatal): {cleanup_error}")