Blame for system_chart_service.py

import tempfile
import os
import shutil
from io import BytesIO
from zipfile import ZipFile
from typing import Dict, Callable, Any, List
from datetime import datetime
from system_chart_extraction import system_chart_extraction

class SystemChartService:
    """
    Service class for System Chart Extraction.
    Handles the orchestration of system chart extraction, visualization creation,
    and result packaging.
    """
    
    def __init__(self):
        pass
    
    async def extract_system_charts(
        self,
        category: str,
        llm,
        pdf_dict: Dict[str, str],
        enable_visualization: bool = False,
        run_system_chart_normalization: bool = False,
        progress_callback: Callable[[Dict[str, Any]], None] = None,
        extraction_id: str = None
    ) -> Dict[str, Any]:
        """
        Main system chart extraction pipeline.
        
        Args:
            category: Database category for extraction
            llm: Language model for extraction
            pdf_dict: Dictionary mapping document titles to full text content
            enable_visualization: Whether to create visualization diagrams
            progress_callback: Function to report progress updates
            
        Returns:
            Dict containing extraction results and metadata
        """
        try:
            # Validate input
            if not pdf_dict:
                return {
                    "success": False,
                    "message": "No files found in the specified category.",
                    "zip_data": None,
                    "filename": None
                }
            
            print(f"📄 Found {len(pdf_dict)} papers to process")
            
            # Initial progress update
            if progress_callback:
                progress_callback({
                    "status": "starting",
                    "progress": 0,
                    "message": "Getting PDF files..."
                })
            
            # Create progress callback function for the extraction process
            def extraction_progress_callback(progress_data):
                if progress_callback:
                    print(f"📊 Progress update: {progress_data}")
                    progress_callback(progress_data)
                    print(f"🔄 Progress updated: {progress_data}")
            
            # Update progress before starting extraction
            if progress_callback:
                progress_callback({
                    "status": "starting",
                    "progress": 5,
                    "message": f"Found {len(pdf_dict)} papers. Starting extraction..."
                })
            
            print(f"🔬 Starting system chart extraction...")
            
            # Create temporary directory for visualization files if needed
            temp_dir = None
            if enable_visualization:
                temp_dir = tempfile.mkdtemp(prefix="system_chart_viz_")
                print(f"📁 Created temporary directory: {temp_dir}")
            
            # Run system chart extraction with progress tracking
            output_system_charts, output_all_details, created_visualizations = await system_chart_extraction(
                pdf_dict, 
                llm, 
                extraction_progress_callback, 
                enable_visualization,
                temp_dir,
                extraction_id,  # Pass extraction_id to save conversations
                run_normalization=run_system_chart_normalization  # ← only normalize when asked/needed
            )
            
            # Create ZIP file in memory
            zip_data = self._create_zip_file(
                output_system_charts, 
                output_all_details, 
                created_visualizations
            )
            
            # Clean up temporary directory
            if temp_dir:
                self._cleanup_temp_directory(temp_dir)
            
            # Prepare result
            viz_message = f" with {len(created_visualizations)} visualizations" if created_visualizations else ""
            
            return {
                "success": True,
                "message": f"System chart extraction completed successfully{viz_message}!",
                "zip_data": zip_data,
                "filename": "system_chart_extraction.zip",
                "charts_count": len(pdf_dict),
                "visualizations_count": len(created_visualizations) if created_visualizations else 0
            }
            
        except Exception as e:
            # Clean up on error
            if 'temp_dir' in locals() and temp_dir:
                self._cleanup_temp_directory(temp_dir)
            
            return {
                "success": False,
                "message": f"Error during system chart extraction: {str(e)}",
                "zip_data": None,
                "filename": None
            }
    
    def _create_zip_file(
        self, 
        output_system_charts: str, 
        output_all_details: str, 
        created_visualizations: List[str]
    ) -> bytes:
        """Create ZIP file containing extraction results and visualizations."""
        zip_buffer = BytesIO()
        
        with ZipFile(zip_buffer, "w") as zipf:
            # Add text results
            zipf.writestr("system_charts.txt", output_system_charts)
            zipf.writestr("all_details.txt", output_all_details)
            
            # Add visualization files if any were created
            if created_visualizations:
                for viz_path in created_visualizations:
                    if os.path.exists(viz_path):
                        # Extract just the filename for the ZIP entry
                        viz_filename = os.path.basename(viz_path)
                        zipf.write(viz_path, f"visualizations/{viz_filename}")
                        print(f"📦 Added visualization to ZIP: {viz_filename}")
        
        # Get ZIP data and close buffer
        zip_data = zip_buffer.getvalue()
        zip_buffer.close()
        
        return zip_data
    
    def _cleanup_temp_directory(self, temp_dir: str) -> None:
        """Clean up temporary directory and its contents."""
        if temp_dir and os.path.exists(temp_dir):
            try:
                shutil.rmtree(temp_dir)
                print(f"🗑️ Cleaned up temporary directory: {temp_dir}")
            except OSError as e:
                print(f"⚠️ Could not clean up temporary directory {temp_dir}: {e}")
                # Don't fail the whole process if cleanup fails
                pass
    
    @staticmethod
    def parse_system_charts_file(file_content: str) -> List[Dict[str, str]]:
        """
        Parse a delimited system charts file into individual chart objects.
        
        Args:
            file_content: Content of the system_charts.txt file
            
        Returns:
            List of dictionaries, each containing:
            - 'chart_number': The chart number (auto-generated)
            - 'title': Paper title
            - 'citation': Citation format (same as title for compatibility)
            - 'content': System chart table content
        """
        charts = []
        lines = file_content.split('\n')
        
        current_chart = None
        chart_number = 1
        
        for i, line in enumerate(lines):
            line = line.strip()
            
            # Check for delimiter line (100 '=' characters)
            if line == '=' * 100:
                # Look for PAPER: line after delimiter
                if i + 1 < len(lines):
                    next_line = lines[i + 1].strip()
                    if next_line.startswith('PAPER:'):
                        # Save previous chart if exists
                        if current_chart:
                            charts.append(current_chart)
                        
                        # Start new chart
                        paper_title = next_line.replace('PAPER:', '').strip()
                        current_chart = {
                            'chart_number': str(chart_number),
                            'title': paper_title,
                            'citation': f'[{paper_title}]',
                            'content': ''
                        }
                        chart_number += 1
                        
            elif line.startswith('SYSTEM_CHART:') and current_chart:
                # Extract the system chart content
                chart_content = line.replace('SYSTEM_CHART:', '').strip()
                current_chart['content'] = chart_content
        
        # Handle last chart if file doesn't end with proper delimiter
        if current_chart:
            charts.append(current_chart)
        
        print(f"📊 Parsed {len(charts)} system charts from file")
        return charts
    
    @staticmethod
    def combine_system_charts_files(file_contents: List[str]) -> str:
        """
        Combine multiple system chart txt files into a single formatted string.
        Preserves the one-to-one relationship between system charts and titles.
        
        Args:
            file_contents: List of system chart file contents to combine
            
        Returns:
            Combined system chart content in the new format
        """
        combined_charts = []
        
        # Parse each file and collect all charts
        for file_content in file_contents:
            charts = SystemChartService.parse_system_charts_file(file_content)
            combined_charts.extend(charts)
        
        # Generate combined output in new format
        combined_output = "COMBINED SYSTEM CHARTS EXTRACTION RESULTS\n"
        combined_output += "=" * 100 + "\n"
        
        for chart in combined_charts:
            delimiter = "=" * 100
            combined_output += f'\n{delimiter}\n'
            combined_output += f'PAPER: {chart["title"]}\n'
            combined_output += f'SYSTEM_CHART: {chart["content"]}\n'
            combined_output += f'{delimiter}\n\n'
        
        # Add summary
        summary_delimiter = "=" * 100
        combined_output += f'\n{summary_delimiter}\n'
        combined_output += f'COMBINATION SUMMARY\n'
        combined_output += f'{summary_delimiter}\n'
        combined_output += f'Total Files Combined: {len(file_contents)}\n'
        combined_output += f'Total System Charts: {len(combined_charts)}\n'
        combined_output += f'Combination Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n'
        
        print(f"🔗 Combined {len(file_contents)} files into {len(combined_charts)} system charts")
        return combined_output
    
    async def combine_multiple_system_chart_files(
        self,
        file_paths: List[str],
        output_filename: str = "combined_system_charts.txt"
    ) -> Dict[str, Any]:
        """
        Combine multiple system chart txt files into a single ZIP file.
        
        Args:
            file_paths: List of file paths to system chart txt files
            output_filename: Name for the output file
            
        Returns:
            Dict containing success status, message, and ZIP data
        """
        try:
            file_contents = []
            
            # Read all files
            for file_path in file_paths:
                if not os.path.exists(file_path):
                    return {
                        "success": False,
                        "message": f"File not found: {file_path}",
                        "zip_data": None,
                        "filename": None
                    }
                
                with open(file_path, 'r', encoding='utf-8') as f:
                    file_contents.append(f.read())
            
            # Combine the files
            combined_content = self.combine_system_charts_files(file_contents)
            
            # Create ZIP file
            zip_buffer = BytesIO()
            with ZipFile(zip_buffer, "w") as zipf:
                zipf.writestr(output_filename, combined_content)
            
            zip_data = zip_buffer.getvalue()
            zip_buffer.close()
            
            return {
                "success": True,
                "message": f"Successfully combined {len(file_paths)} system chart files!",
                "zip_data": zip_data,
                "filename": "combined_system_charts.zip",
                "files_combined": len(file_paths)
            }
            
        except Exception as e:
            return {
                "success": False,
                "message": f"Error combining system chart files: {str(e)}",
                "zip_data": None,
                "filename": None
            }
    
    @staticmethod
    def get_chart_citations(file_content: str) -> List[str]:
        """
        Extract just the citations from a system charts file.
        
        Args:
            file_content: Content of the system_charts.txt file
            
        Returns:
            List of citation strings
        """
        charts = SystemChartService.parse_system_charts_file(file_content)
        return [chart.get('citation', chart.get('title', f"Chart {chart['chart_number']}")) 
                for chart in charts if chart.get('chart_number')]