"""
Configuration utilities for loading and managing settings.
"""
import yaml
import os
import shutil
from datetime import datetime
from typing import Dict, Tuple


def load_config(config_path: str = "config.yaml") -> Dict:
    """Load configuration from YAML file."""
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Config file not found: {config_path}")

    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    return config


def create_directories(config: Dict):
    """Create necessary directories for checkpoints and logs."""
    checkpoint_dir = config["training"]["checkpoint_dir"]
    log_dir = config["training"]["log_dir"]

    os.makedirs(checkpoint_dir, exist_ok=True)
    os.makedirs(log_dir, exist_ok=True)


def create_run_directory(config: Dict, config_path: str = "config.yaml") -> Tuple[str, str, str]:
    """
    Create a timestamped run directory for this training session.
    Latest run is saved to 'latest/', previous runs are moved to 'runs/'.

    Args:
        config: Configuration dictionary
        config_path: Path to the config file

    Returns:
        Tuple of (run_dir, checkpoint_dir, log_dir)
    """
    # Create timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # Check if 'latest' directory exists
    latest_dir = "latest"
    if os.path.exists(latest_dir):
        # Move existing 'latest' to 'runs' with timestamp from its config
        try:
            # Try to read timestamp from existing latest directory
            old_config_path = os.path.join(latest_dir, "config.yaml")
            if os.path.exists(old_config_path):
                # Get modification time of the old config
                mtime = os.path.getmtime(old_config_path)
                old_timestamp = datetime.fromtimestamp(mtime).strftime("%Y%m%d_%H%M%S")
            else:
                old_timestamp = "unknown"

            # Move to runs directory
            os.makedirs("runs", exist_ok=True)
            archive_dir = os.path.join("runs", f"run_{old_timestamp}")
            shutil.move(latest_dir, archive_dir)

            # Clean up timestamp files in archived directory
            for file in os.listdir(archive_dir):
                if file.startswith("timestamp_") and file.endswith(".txt"):
                    timestamp_file_path = os.path.join(archive_dir, file)
                    os.remove(timestamp_file_path)

            print(f"Archived previous run to: {archive_dir}")
        except Exception as e:
            print(f"Warning: Could not archive previous run: {e}")

    # Create new 'latest' directory
    run_dir = latest_dir
    os.makedirs(run_dir, exist_ok=True)

    # Create subdirectories
    checkpoint_dir = os.path.join(run_dir, "checkpoints")
    log_dir = os.path.join(run_dir, "logs")
    os.makedirs(checkpoint_dir, exist_ok=True)
    os.makedirs(log_dir, exist_ok=True)

    # Copy config file to run directory with timestamp
    config_copy_path = os.path.join(run_dir, "config.yaml")
    shutil.copy(config_path, config_copy_path)

    # Also save a timestamped copy for reference
    timestamp_file = os.path.join(run_dir, f"timestamp_{timestamp}.txt")
    with open(timestamp_file, "w") as f:
        f.write(f"Training started at: {timestamp}\n")

    print(f"Created run directory: {run_dir}")
    print(f"Training timestamp: {timestamp}")
    print(f"Config saved to: {config_copy_path}")

    return run_dir, checkpoint_dir, log_dir