csv_handler.py

"""
CSV handling functionality for Animagine Prompt Node
"""

import pandas as pd
from typing import Optional, Dict
import os
import time
from .logger import logger

class CSVHandler:
    """
    Handles loading and processing of CSV files containing character data
    """
    
    REQUIRED_COLUMNS = ['GENDER', 'CHARACTER', 'COPYRIGHT']
    
    def __init__(self):
        self.current_data = None
        self.cache = {}
    
    def validate_csv_structure(self, filepath: str) -> bool:
        """
        Validates if the CSV file has the required structure
        """
        try:
            if not os.path.exists(filepath):
                logger.log_error(
                    FileNotFoundError(f"CSV file not found: {filepath}"),
                    "CSV Validation",
                    {"filepath": filepath}
                )
                return False
                
            df = pd.read_csv(filepath)
            logger.log_csv_operation(
                "Validation",
                filepath,
                {"columns_found": df.columns.tolist()}
            )
            
            missing_columns = [col for col in self.REQUIRED_COLUMNS if col not in df.columns]
            if missing_columns:
                logger.log_error(
                    ValueError("Missing required columns"),
                    "CSV Validation",
                    {"missing_columns": missing_columns}
                )
                return False
                
            return True
            
        except Exception as e:
            logger.log_error(e, "CSV Validation", {"filepath": filepath})
            return False
    
    def load_csv(self, filepath: str) -> Optional[pd.DataFrame]:
        """
        Loads CSV file with basic caching
        """
        try:
            start_time = time.time()
            
            if filepath in self.cache:
                logger.log_csv_operation(
                    "Cache Hit",
                    filepath,
                    {"cache_size": len(self.cache)}
                )
                return self.cache[filepath]
            
            if not self.validate_csv_structure(filepath):
                return None
            
            df = pd.read_csv(filepath)
            self.cache[filepath] = df
            self.current_data = df
            
            end_time = time.time()
            logger.log_performance(
                "CSV Load",
                start_time,
                end_time,
                {
                    "filepath": filepath,
                    "rows": len(df),
                    "cached": True
                }
            )
            
            return df
            
        except Exception as e:
            logger.log_error(e, "CSV Load", {"filepath": filepath})
            return None
    
    def get_entry(self, index: int) -> Optional[Dict[str, str]]:
        """
        Gets a specific entry from the loaded CSV
        """
        try:
            if self.current_data is None:
                logger.log_error(
                    ValueError("No CSV data loaded"),
                    "Get Entry",
                    {"index": index}
                )
                return None
            
            if index >= len(self.current_data):
                logger.log_error(
                    IndexError(f"Index {index} out of bounds"),
                    "Get Entry",
                    {
                        "index": index,
                        "max_index": len(self.current_data) - 1
                    }
                )
                return None
                
            row = self.current_data.iloc[index]
            entry = {
                'gender': row['GENDER'],
                'character': row['CHARACTER'],
                'copyright': row['COPYRIGHT']
            }
            
            logger.log_csv_operation(
                "Get Entry",
                None,
                {"index": index, "entry": entry}
            )
            
            return entry
            
        except Exception as e:
            logger.log_error(e, "Get Entry", {"index": index})
            return None