file_processors.py

import base64
import json
from typing import Dict, Any
import PyPDF2
from docx import Document
import pandas as pd
import nbformat

class BaseProcessor:
    """Base class for all file processors"""
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        raise NotImplementedError("Subclasses must implement process method")

class PDFProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            pages = []
            for n, page in enumerate(pdf_reader.pages):
                pages.append({"page": n, "content": page.extract_text()})
            return {
                "content": pages,
                "metadata": {
                    "total_pages": len(pages),
                    "file_type": "PDF"
                }
            }

class DocxProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        doc = Document(file_path)
        paragraphs = []
        
        for n, paragraph in enumerate(doc.paragraphs, 1):
            if paragraph.text.strip():
                paragraphs.append({
                    "paragraph": n,
                    "content": paragraph.text
                })
        
        return {
            "content": paragraphs,
            "metadata": {
                "total_paragraphs": len(paragraphs),
                "file_type": "DOCX"
            }
        }

class ImageProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        try:
            with open(file_path, 'rb') as file:
                content = base64.b64encode(file.read()).decode('ascii')
                return {
                    "content": content,
                    "metadata": {
                        "file_type": "IMAGE",
                        "encoding": "base64"
                    }
                }
        except Exception as e:
            raise Exception(f"Error processing image: {str(e)}")

class JSONProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            try:
                data = json.load(file)
                return {
                    "content": data,
                    "metadata": {
                        "file_type": "JSON",
                        "is_valid": True
                    }
                }
            except json.JSONDecodeError as e:
                return {
                    "content": None,
                    "metadata": {
                        "file_type": "JSON",
                        "is_valid": False,
                        "error": str(e)
                    }
                }

class TextProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "TEXT",
                    "encoding": "utf-8"
                }
            }

class CSVProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        try:
            df = pd.read_csv(file_path)
            return {
                "content": df.to_dict(orient='records'),
                "metadata": {
                    "file_type": "CSV",
                    "rows": len(df),
                    "columns": len(df.columns),
                    "column_names": list(df.columns),
                    "summary": {
                        "first_few_rows": df.head().to_dict(orient='records'),
                        "statistics": df.describe().to_dict()
                    }
                }
            }
        except Exception as e:
            return {
                "content": None,
                "metadata": {
                    "file_type": "CSV",
                    "error": str(e)
                }
            }

class ExcelProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        try:
            # Get sheet names first
            xl = pd.ExcelFile(file_path)
            sheets = xl.sheet_names
            
            # Read the first sheet by default
            df = pd.read_excel(file_path, sheet_name=sheets[0])
            
            # Read all sheets
            all_sheets = {}
            for sheet in sheets:
                all_sheets[sheet] = pd.read_excel(file_path, sheet_name=sheet).to_dict(orient='records')
            
            return {
                "content": all_sheets,
                "metadata": {
                    "file_type": "EXCEL",
                    "sheets": sheets,
                    "current_sheet": {
                        "name": sheets[0],
                        "rows": len(df),
                        "columns": len(df.columns),
                        "column_names": list(df.columns),
                        "summary": {
                            "first_few_rows": df.head().to_dict(orient='records'),
                            "statistics": df.describe().to_dict()
                        }
                    }
                }
            }
        except Exception as e:
            return {
                "content": None,
                "metadata": {
                    "file_type": "EXCEL",
                    "error": str(e)
                }
            }

class PythonProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "PYTHON",
                    "encoding": "utf-8"
                }
            }

class JupyterNotebookProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        try:
            with open(file_path, 'r') as file:
                nb = nbformat.read(file, as_version=4)
                content = []
                for cell in nb.cells:
                    cell_info = {
                        "type": cell.cell_type,
                        "content": cell.source
                    }
                    if cell.cell_type == "code" and cell.outputs:
                        cell_info["outputs"] = [str(output) for output in cell.outputs]
                    content.append(cell_info)
                
                return {
                    "content": content,
                    "metadata": {
                        "file_type": "JUPYTER_NOTEBOOK",
                        "total_cells": len(content),
                        "cell_types": list(set(cell["type"] for cell in content))
                    }
                }
        except Exception as e:
            return {
                "content": None,
                "metadata": {
                    "file_type": "JUPYTER_NOTEBOOK",
                    "error": str(e)
                }
            }

class SVGProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        try:
            with open(file_path, 'r') as file:
                return {
                    "content": file.read(),
                    "metadata": {
                        "file_type": "SVG",
                        "encoding": "utf-8"
                    }
                }
        except Exception as e:
            return {
                "content": None,
                "metadata": {
                    "file_type": "SVG",
                    "error": str(e)
                }
            }

class JavaScriptProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "JAVASCRIPT",
                    "encoding": "utf-8"
                }
            }

class HTMLProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "HTML",
                    "encoding": "utf-8"
                }
            }

class CSSProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "CSS",
                    "encoding": "utf-8"
                }
            }

class JavaProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "JAVA",
                    "encoding": "utf-8"
                }
            }

class CppProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "CPP",
                    "encoding": "utf-8"
                }
            }

class HeaderProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "HEADER",
                    "encoding": "utf-8"
                }
            }

class ShellScriptProcessor(BaseProcessor):
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            return {
                "content": content,
                "metadata": {
                    "file_type": "SHELL_SCRIPT",
                    "encoding": "utf-8"
                }
            }

class CodeFileProcessor(BaseProcessor):
    """Universal processor for code files that just need to be read as text"""
    @classmethod
    def process(cls, file_path: str) -> Dict[str, Any]:
        with open(file_path, 'r') as file:
            content = file.read()
            ext = file_path[file_path.rfind('.'):].upper()[1:]  # Remove the dot and capitalize
            return {
                "content": content,
                "metadata": {
                    "file_type": ext,
                    "encoding": "utf-8"
                }
            }

class FileProcessorFactory:
    """Factory class to get the appropriate processor for a file type"""
    _processors = {
        '.pdf': PDFProcessor,
        '.docx': DocxProcessor,
        '.png': ImageProcessor,
        '.jpg': ImageProcessor,
        '.jpeg': ImageProcessor,
        '.json': JSONProcessor,
        '.txt': TextProcessor,
        '.csv': CSVProcessor,
        '.xlsx': ExcelProcessor,
        '.xls': ExcelProcessor,
        '.py': PythonProcessor,
        '.ipynb': JupyterNotebookProcessor,
        '.svg': SVGProcessor,
        # Universal code file processor for various file types
        '.js': CodeFileProcessor,
        '.html': CodeFileProcessor,
        '.htm': CodeFileProcessor,
        '.css': CodeFileProcessor,
        '.java': CodeFileProcessor,
        '.cpp': CodeFileProcessor,
        '.cc': CodeFileProcessor,
        '.cxx': CodeFileProcessor,
        '.h': CodeFileProcessor,
        '.hpp': CodeFileProcessor,
        '.sh': CodeFileProcessor,
        '.bash': CodeFileProcessor
    }
    
    @classmethod
    def get_processor(cls, file_path: str) -> BaseProcessor:
        ext = file_path[file_path.rfind('.'):].lower()
        processor = cls._processors.get(ext)
        if not processor:
            raise ValueError(f"Unsupported file type: {ext}")
        return processor