from typing import Dict, List, Optional
from pathlib import Path
import pandas as pd
from llama_index.readers.file import CSVReader
from llama_index.schema import Document

class EnhancedCSVReader:
    """Enhanced CSV reader with metadata extraction capabilities."""
    
    def __init__(self):
        self.csv_reader = CSVReader()
        
    def load_data(self, file_path: str) -> List[Document]:
        """Load CSV file and extract documents with metadata."""
        # Load the CSV file
        documents = self.csv_reader.load_data(file_path)
        
        # Extract and add metadata
        csv_metadata = self._extract_metadata(file_path)
        
        # Enhance documents with metadata
        for doc in documents:
            doc.metadata.update(csv_metadata)
            
        return documents
    
    def _extract_metadata(self, file_path: str) -> Dict:
        """Extract useful metadata from CSV file."""
        df = pd.read_csv(file_path)
        filename = Path(file_path).name
        
        # Extract column information
        columns = df.columns.tolist()
        dtypes = {col: str(df[col].dtype) for col in columns}
        
        # Extract sample values (first 3 non-null values per column)
        samples = {}
        for col in columns:
            non_null_values = df[col].dropna().head(3).tolist()
            samples[col] = [str(val) for val in non_null_values]
            
        # Basic statistics
        row_count = len(df)
        
        return {
            "filename": filename,
            "columns": columns,
            "dtypes": dtypes,
            "samples": samples,
            "row_count": row_count
        }