import logging
import csv
import os
from datetime import datetime
from typing import Dict, Any
from pprint import pprint


def configure_logger(level: int = logging.INFO, name: str = "agentarium") -> logging.Logger:
    logger = logging.getLogger(name)
    logger.handlers.clear()  # clear existing handlers

    # log to stdout
    handler = logging.StreamHandler()
    handler.setLevel(level)
    formatter = logging.Formatter("%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(level)
    logger.propagate = False  # prevent double logging
    return logger


class CSVMetricsLogger:
    """Handles logging of training metrics to CSV files."""
    
    def __init__(self, log_dir: str, experiment_name: str):
        """
        Initialize CSV metrics logger.
        
        Args:
            log_dir: Directory to store CSV files
            experiment_name: Name of the experiment for the CSV filename
        """
        self.log_dir = log_dir
        self.experiment_name = experiment_name
        self.csv_filepath = None
        self.csv_headers_written = False
        self.csv_fieldnames = None
        
        # Create log directory and CSV file path
        os.makedirs(log_dir, exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_filename = f"{experiment_name}_metrics_{timestamp}.csv"
        self.csv_filepath = os.path.join(log_dir, csv_filename)
    
    def log_metrics(self, metrics: Dict[str, Any], step: int) -> None:
        """
        Log metrics to CSV file.
        
        Args:
            metrics: Dictionary of metrics to log
            step: Training step number
        """
        try:
            # Flatten nested metrics dictionary
            flat_metrics = self._flatten_metrics(metrics)
            
            # Add step information
            flat_metrics['step'] = step
            
            # Check if file exists and if headers need to be written
            file_exists = os.path.exists(self.csv_filepath)
            
            with open(self.csv_filepath, 'a', newline='') as csvfile:
                if not file_exists or not self.csv_headers_written:
                    # Write headers
                    fieldnames = sorted(flat_metrics.keys())
                    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                    writer.writeheader()
                    self.csv_headers_written = True
                    self.csv_fieldnames = fieldnames
                else:
                    # Use existing fieldnames, add new ones if necessary
                    new_fields = set(flat_metrics.keys()) - set(self.csv_fieldnames)
                    if new_fields:
                        # If new fields are found, we need to rewrite the file with new headers
                        self._rewrite_csv_with_new_headers(flat_metrics)
                        return
                    fieldnames = self.csv_fieldnames
                
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                # Fill missing fields with empty values
                row_data = {field: flat_metrics.get(field, '') for field in fieldnames}
                writer.writerow(row_data)
                
        except Exception as e:
            pprint(f"Warning: Failed to log metrics to CSV: {e}")
    
    def _flatten_metrics(self, metrics: Dict[str, Any]) -> Dict[str, Any]:
        """
        Flatten nested metrics dictionary.
        
        Args:
            metrics: Dictionary that may contain nested dictionaries
            
        Returns:
            Flattened dictionary with keys like 'parent/child'
        """
        flat_metrics = {}
        for key, value in metrics.items():
            if isinstance(value, (int, float)):
                flat_metrics[key] = value
            elif isinstance(value, dict):
                for sub_key, sub_value in value.items():
                    if isinstance(sub_value, (int, float)):
                        flat_metrics[f"{key}/{sub_key}"] = sub_value
            else:
                # Convert other types to string
                flat_metrics[key] = str(value)
        return flat_metrics
    
    def _rewrite_csv_with_new_headers(self, new_metrics: Dict[str, Any]) -> None:
        """
        Rewrite CSV file when new metric fields are discovered.
        
        Args:
            new_metrics: New metrics dictionary containing new fields
        """
        try:
            # Read existing data
            existing_data = []
            if os.path.exists(self.csv_filepath):
                with open(self.csv_filepath, 'r') as csvfile:
                    reader = csv.DictReader(csvfile)
                    existing_data = list(reader)
            
            # Determine all fieldnames
            all_fieldnames = set(new_metrics.keys())
            for row in existing_data:
                all_fieldnames.update(row.keys())
            all_fieldnames = sorted(all_fieldnames)
            
            # Rewrite file
            with open(self.csv_filepath, 'w', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=all_fieldnames)
                writer.writeheader()
                
                # Write existing data
                for row in existing_data:
                    padded_row = {field: row.get(field, '') for field in all_fieldnames}
                    writer.writerow(padded_row)
                
                # Write new data
                new_row = {field: new_metrics.get(field, '') for field in all_fieldnames}
                writer.writerow(new_row)
            
            self.csv_fieldnames = all_fieldnames
            
        except Exception as e:
            pprint(f"Warning: Failed to rewrite CSV with new headers: {e}")
    
    @property
    def filepath(self) -> str:
        """Get the CSV file path."""
        return self.csv_filepath
