"""
IO Manager for handling file paths and saving artifacts.

Ref: revisions.txt Section 1
"""

import os
import json
import logging
import pandas as pd
from typing import Any, Dict, Optional, List
import matplotlib.pyplot as plt

logger = logging.getLogger(__name__)

class IOManager:
    """
    Manages input/output operations, including path construction and saving various artifacts.
    """

    def __init__(self, cfg: Dict[str, Any], run_id: str):
        """
        Initialize the IOManager.

        Args:
            cfg (Dict[str, Any]): The pipeline configuration dictionary.
            run_id (str): The unique identifier for the current run.
        """
        self.cfg = cfg
        self.run_id = run_id
        
        # Extract base directories, providing defaults if missing
        self.base_dirs = cfg.get('base_dirs', {})
        self.results_dir = self._resolve_path(self.base_dirs.get('results', 'results'))
        self.models_dir = self._resolve_path(self.base_dirs.get('models', 'models'))
        self.logs_dir = self._resolve_path(self.base_dirs.get('logs', 'logs'))
        
        # Specific directories for the current run
        self.run_results_dir = os.path.join(self.results_dir, self.run_id)
        self.run_models_dir = os.path.join(self.models_dir, self.run_id)
        self.run_logs_dir = os.path.join(self.logs_dir, self.run_id)
        self.run_figures_dir = os.path.join(self.run_results_dir, 'figures') # Figures within results
        
        # Create directories if they don't exist
        os.makedirs(self.run_results_dir, exist_ok=True)
        os.makedirs(self.run_models_dir, exist_ok=True)
        os.makedirs(self.run_logs_dir, exist_ok=True)
        os.makedirs(self.run_figures_dir, exist_ok=True)
        
        logger.info(f"IOManager initialized for run {self.run_id}.")
        logger.info(f"  Results Dir: {self.run_results_dir}")
        logger.info(f"  Models Dir:  {self.run_models_dir}")
        logger.info(f"  Logs Dir:    {self.run_logs_dir}")
        logger.info(f"  Figures Dir: {self.run_figures_dir}")

    def _resolve_path(self, path: str) -> str:
        """
        Resolves a path relative to the project root.
        Assumes this file is in src/ for relative path calculation.
        """
        if os.path.isabs(path):
            return path
        else:
            # Assumes src/io_manager.py structure
            script_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(script_dir) 
            # Go up one level from src to get to package root
            package_root = os.path.dirname(project_root)
            return os.path.abspath(os.path.join(package_root, path))
            
    def path(self, section: str, name: str, suffix: Optional[str] = None) -> str:
        """
        Constructs a full path for an artifact within a specific run section.

        Args:
            section (str): The base directory section ('results', 'models', 'logs', 'figures').
            name (str): The base name of the file (without extension or run_id typically).
            suffix (Optional[str]): File extension (e.g., '.txt', '.png'). Auto-added by save methods.

        Returns:
            str: The full, absolute path to the artifact.
                 Includes the run_id in the path structure.
        """
        base_path = ""
        if section == 'results':
            base_path = self.run_results_dir
        elif section == 'models':
            base_path = self.run_models_dir
        elif section == 'logs':
            base_path = self.run_logs_dir
        elif section == 'figures':
             base_path = self.run_figures_dir
        else:
            raise ValueError(f"Unknown path section: '{section}'. Must be one of results, models, logs, figures.")
            
        filename = name
        if suffix:
             if not suffix.startswith('.'):
                  suffix = '.' + suffix
             filename += suffix
             
        full_path = os.path.join(base_path, filename)
        return full_path

    def get_fold_dirs(self, fold_num: int) -> Dict[str, str]:
        """
        Creates and returns paths for fold-specific subdirectories.

        Args:
            fold_num (int): The fold number (1-based).

        Returns:
            Dict[str, str]: Dictionary mapping section names to fold-specific directory paths.
        """
        fold_suffix = f"fold_{fold_num}"
        fold_dirs = {}

        # Define paths for fold-specific subdirectories
        fold_dirs['results'] = os.path.join(self.run_results_dir, fold_suffix)
        fold_dirs['models'] = os.path.join(self.run_models_dir, fold_suffix)
        fold_dirs['logs'] = os.path.join(self.run_logs_dir, fold_suffix) # Optional: logs per fold?
        fold_dirs['figures'] = os.path.join(self.run_figures_dir, fold_suffix)

        # Create these directories
        for section, path in fold_dirs.items():
            try:
                os.makedirs(path, exist_ok=True)
                logger.debug(f"Ensured fold directory exists: {path}")
            except OSError as e:
                logger.error(f"Failed to create fold directory '{path}' for section '{section}': {e}")
                # Decide how to handle error: raise, return partial dict, etc.
                # For now, let's log and continue, returning potentially incomplete dict

        logger.info(f"Fold {fold_num} directories prepared:")
        for section, path in fold_dirs.items():
             logger.info(f"  {section.capitalize()}: {path}")

        return fold_dirs

    # --- Save Methods (Task 1.1) --- #
    def save_json(self, data: Dict[str, Any], name: str, section: str = 'results', indent: int = 4, use_txt: bool = False):
        """
        Saves dictionary data to a JSON file (or .txt if specified) in the target section.
        """
        suffix = '.txt' if use_txt else '.json'
        file_path = self.path(section, name, suffix=suffix)
        logger.info(f"Saving JSON data to {file_path}...")
        try:
            os.makedirs(os.path.dirname(file_path), exist_ok=True) 
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=indent)
            logger.debug(f"Successfully saved JSON to {file_path}")
        except TypeError as e:
             logger.error(f"TypeError saving JSON to {file_path}. Data may contain non-serializable types: {e}")
        except Exception as e:
            logger.error(f"Failed to save JSON to {file_path}: {e}", exc_info=True)
            
    def save_df(self, df: pd.DataFrame, name: str, section: str = 'results'):
        """
        Saves DataFrame to CSV or Parquet based on estimated size and config.
        Uses 'output.dataframe_save_format' and 'output.dataframe_max_csv_mb' from config.
        """
        if df is None or df.empty:
             logger.warning(f"Attempted to save empty DataFrame '{name}'. Skipping.")
             return None # Return None to indicate skip

        output_cfg = self.cfg.get('output', {})
        save_format_pref = output_cfg.get('dataframe_save_format', 'parquet_if_large') # Default pref
        max_csv_mb = output_cfg.get('dataframe_max_csv_mb', 100) # Default size limit

        file_path = None
        saved_format = None

        try:
            size_mb = df.memory_usage(index=True, deep=True).sum() / (1024**2)
            logger.debug(f"DataFrame '{name}' estimated size: {size_mb:.2f} MB (Max CSV Size: {max_csv_mb} MB)")

            use_parquet = False
            if save_format_pref == 'parquet':
                 use_parquet = True
            elif save_format_pref == 'parquet_if_large' and size_mb > max_csv_mb:
                 use_parquet = True
            # else: save_format_pref == 'csv' or (parquet_if_large and size_mb <= max_csv_mb) -> use_csv

            if use_parquet:
                try:
                    # Check for pyarrow installation before attempting Parquet save
                    import pyarrow # noqa: F401
                    suffix = '.parquet'
                    saved_format = 'Parquet'
                    file_path = self.path(section, name, suffix=suffix)
                    logger.info(f"Saving DataFrame '{name}' as {saved_format} to {file_path}...")
                    os.makedirs(os.path.dirname(file_path), exist_ok=True)
                    df.to_parquet(file_path, index=True)
                except ImportError:
                    logger.warning(f"Cannot save DataFrame '{name}' as Parquet. Missing 'pyarrow'. Falling back to CSV.")
                    use_parquet = False # Force CSV save
                except Exception as parquet_e:
                    logger.error(f"Failed to save DataFrame '{name}' as Parquet: {parquet_e}", exc_info=True)
                    # Optionally fallback? For now, let error propagate if it's not ImportError
                    raise parquet_e

            if not use_parquet: # Save as CSV if format is 'csv', fallback, or size threshold not met
                suffix = '.csv'
                saved_format = 'CSV'
                file_path = self.path(section, name, suffix=suffix)
                logger.info(f"Saving DataFrame '{name}' as {saved_format} to {file_path}...")
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                df.to_csv(file_path, index=True)

            logger.debug(f"Successfully saved DataFrame to {file_path}")
            return file_path # Return path where it was saved

        except Exception as e:
            logger.error(f"Failed to save DataFrame '{name}' (intended format: {saved_format or 'N/A'}): {e}", exc_info=True)
            return None # Return None on failure
        
    def save_figure(self, fig: plt.Figure, name: str, section: str = 'figures', **kwargs):
        """
        Saves matplotlib figure using config settings (dpi).
        """
        file_path = self.path(section, name, suffix='.png')
        logger.info(f"Saving figure to {file_path}...")
        try:
            os.makedirs(os.path.dirname(file_path), exist_ok=True) 
            output_cfg = self.cfg.get('output', {})
            dpi = kwargs.pop('dpi', output_cfg.get('figure_dpi', 150))
            
            if hasattr(fig, 'tight_layout'):
                try:
                    fig.tight_layout()
                except Exception as tl_e:
                     logger.warning(f"Could not apply tight_layout to figure '{name}': {tl_e}")

            fig.savefig(file_path, dpi=dpi, bbox_inches='tight', **kwargs)
            logger.debug(f"Successfully saved figure to {file_path}")
        except Exception as e:
             logger.error(f"Failed to save figure to {file_path}: {e}", exc_info=True)
        finally:
             plt.close(fig) # Close figure to free memory

# Example Usage
if __name__ == '__main__':
    # Mock config for testing
    mock_config = {
        'base_dirs': {'results': 'temp_results', 'models': 'temp_models', 'logs': 'temp_logs'},
        'output': {'figure_dpi': 120}
    }
    mock_run_id = "20250418_110000_testabc"
    
    # Create mock directories for test
    if not os.path.exists('temp_results'): os.makedirs('temp_results')
    if not os.path.exists('temp_models'): os.makedirs('temp_models')
    if not os.path.exists('temp_logs'): os.makedirs('temp_logs')

    io = IOManager(mock_config, mock_run_id)
    
    print(f"Results Path: {io.path('results', 'metrics', '.txt')}")
    print(f"Models Path: {io.path('models', 'gru_model', '.keras')}")
    print(f"Figures Path: {io.path('figures', 'calibration_plot', '.png')}")
    print(f"Logs Path: {io.path('logs', 'pipeline_log', '.log')}")

    # Test saving
    test_dict = {'a': 1, 'b': [2, 3], 'c': 'test'}
    io.save_json(test_dict, 'test_data', section='results')
    io.save_json(test_dict, 'report_data', section='results', use_txt=True)

    test_df_small = pd.DataFrame(np.random.randn(100, 5), columns=list('ABCDE'))
    io.save_df(test_df_small, 'small_data', section='results')

    # Clean up mock directories
    import shutil
    if os.path.exists('temp_results'): shutil.rmtree('temp_results')
    if os.path.exists('temp_models'): shutil.rmtree('temp_models')
    if os.path.exists('temp_logs'): shutil.rmtree('temp_logs')
 