"""Document loaders for different file formats."""
from typing import Dict, Any, Union, TextIO
import pandas as pd
import json
from PyPDF2 import PdfReader
import logging
import io

logger = logging.getLogger(__name__)

class DocumentLoader:
    """Base class for document loaders with enhanced format handling"""

    @staticmethod
    def load_pdf(file_path: str) -> str:
        """Load and extract text from PDF file"""
        try:
            reader = PdfReader(file_path)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text.strip()
        except Exception as e:
            logger.error(f"Error loading PDF: {str(e)}")
            raise

    @staticmethod
    def load_csv(file_path: str) -> str:
        """Load and format CSV data as structured text"""
        try:
            df = pd.read_csv(file_path)
            # Convert DataFrame to a more natural text format
            text_content = []

            # Add header description
            text_content.append(f"CSV Document with {len(df)} records and columns: {', '.join(df.columns)}\n")

            # Convert each row to a descriptive text
            for _, row in df.iterrows():
                row_text = "Record:\n"
                for col in df.columns:
                    row_text += f"- {col}: {row[col]}\n"
                text_content.append(row_text)

            return "\n".join(text_content)
        except Exception as e:
            logger.error(f"Error loading CSV: {str(e)}")
            raise

    @staticmethod
    def load_excel(file_path: str) -> str:
        """Load and format Excel data as structured text"""
        try:
            df = pd.read_excel(file_path)
            text_content = []

            # Add header description
            text_content.append(f"Excel Document with {len(df)} records and columns: {', '.join(df.columns)}\n")

            # Convert each row to a descriptive text
            for _, row in df.iterrows():
                row_text = "Record:\n"
                for col in df.columns:
                    row_text += f"- {col}: {row[col]}\n"
                text_content.append(row_text)

            return "\n".join(text_content)
        except Exception as e:
            logger.error(f"Error loading Excel: {str(e)}")
            raise

    @staticmethod
    def load_json(file_path: str) -> str:
        """Load and format JSON data as structured text"""
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)

            def process_json(obj, prefix="", depth=0):
                """Recursively process JSON structure into readable text"""
                if depth > 10:  # Prevent infinite recursion
                    return "Maximum depth reached"

                text_parts = []
                indent = "  " * depth

                if isinstance(obj, dict):
                    for key, value in obj.items():
                        if isinstance(value, (dict, list)):
                            text_parts.append(f"{indent}{prefix}{key}:")
                            text_parts.append(process_json(value, "", depth + 1))
                        else:
                            text_parts.append(f"{indent}{prefix}{key}: {value}")
                elif isinstance(obj, list):
                    for i, item in enumerate(obj):
                        if isinstance(item, (dict, list)):
                            text_parts.append(f"{indent}Item {i + 1}:")
                            text_parts.append(process_json(item, "", depth + 1))
                        else:
                            text_parts.append(f"{indent}- {item}")
                else:
                    text_parts.append(f"{indent}{prefix}{obj}")

                return "\n".join(text_parts)

            return process_json(data)

        except Exception as e:
            logger.error(f"Error loading JSON: {str(e)}")
            raise

    @staticmethod
    def load_text(file_path: str) -> str:
        """Load plain text file"""
        try:
            with open(file_path, 'r') as f:
                return f.read().strip()
        except Exception as e:
            logger.error(f"Error loading text file: {str(e)}")
            raise

    @staticmethod
    def detect_file_type(file_path: str) -> str:
        """Detect file type from extension"""
        lower_path = file_path.lower()
        if lower_path.endswith('.pdf'):
            return 'pdf'
        elif lower_path.endswith('.csv'):
            return 'csv'
        elif lower_path.endswith(('.xls', '.xlsx')):
            return 'excel'
        elif lower_path.endswith('.json'):
            return 'json'
        elif lower_path.endswith('.txt'):
            return 'text'
        else:
            return 'unknown'

    @classmethod
    def load(cls, file_path: str) -> str:
        """Load document based on file type"""
        file_type = cls.detect_file_type(file_path)
        logger.info(f"Loading document of type: {file_type} from {file_path}")

        loader_map = {
            'pdf': cls.load_pdf,
            'csv': cls.load_csv,
            'excel': cls.load_excel,
            'json': cls.load_json,
            'text': cls.load_text
        }

        if file_type in loader_map:
            content = loader_map[file_type](file_path)
            logger.info(f"Successfully loaded {file_type} document, content length: {len(content)}")
            return content
        else:
            raise ValueError(f"Unsupported file type: {file_type}")