#!/usr/bin/env python3

import os
import subprocess
import toml
import argparse
from datetime import datetime

# Path to the configuration file
CONFIG_PATH = "/etc/snap-slack/config.toml"

# Read configuration from a TOML file
def read_config(file_path=CONFIG_PATH):
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Configuration file not found: {file_path}")
    
    with open(file_path, 'r') as f:
        config = toml.load(f)
    return config

# Create a new snapshot with the current date and time
def create(config):
    date_str = datetime.now().strftime("%Y%m%d-%H%M%S")
    snapshot_name = f"{config['snapshot']['snapshot_prefix']}{date_str}"
    snapshot_path = os.path.join(config['snapshot']['snapshot_dir'], snapshot_name)

    # Run the Btrfs snapshot command
    cmd = ["btrfs", "subvolume", "snapshot",
           os.path.join(config['snapshot']['btrfs_mount_point'], config['snapshot']['root_subvolume']),
           snapshot_path]
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    if result.returncode == 0:
        print(f"Created snapshot: {snapshot_name}")
        add_elilo_entry(config, snapshot_name)
        return snapshot_name
    else:
        raise Exception(f"Failed to create snapshot: {result.stderr.decode()}")

# Add an entry for the snapshot in elilo.conf
def add_elilo_entry(config, snapshot):
    # Get the extra boot options from the config, if any
    extra_boot_options = config['elilo'].get('extra_boot_options', '')

    # Construct the boot entry with optional extra boot options
    entry = f"""
image=vmlinuz
        label={snapshot}
        root={config['elilo']['root_partition']}
        append="rootflags=subvol={config['snapshot']['snapshot_dir']}/{snapshot} ro {extra_boot_options}"
        read-only
        initrd=initrd.gz
"""

    # Append to elilo.conf if the entry doesn't already exist
    elilo_conf = config['elilo']['elilo_conf']
    with open(elilo_conf, 'r+') as f:
        content = f.read()
        if snapshot not in content:
            print(f"Adding entry for snapshot: {snapshot}")
            f.write(entry)

# Remove an entry for a deleted snapshot from elilo.conf
def remove_elilo_entry(config, snapshot):
    elilo_conf = config['elilo']['elilo_conf']
    with open(elilo_conf, 'r') as f:
        lines = f.readlines()

    with open(elilo_conf, 'w') as f:
        for line in lines:
            if snapshot not in line:
                f.write(line)

    print(f"Removed entry for snapshot: {snapshot}")

# List all current snapshots
def list_snapshots(config):
    snapshots = []
    for entry in os.listdir(config['snapshot']['snapshot_dir']):
        entry_path = os.path.join(config['snapshot']['snapshot_dir'], entry)
        if os.path.isdir(entry_path) and entry.startswith(config['snapshot']['snapshot_prefix']):
            snapshots.append(entry)
    return snapshots

# Remove snapshots older than the retention period
def remove_old_snapshots(config):
    retention_days = config['snapshot']['retention_days']
    now = datetime.now()

    for snapshot in list_snapshots(config):
        snapshot_path = os.path.join(config['snapshot']['snapshot_dir'], snapshot)
        created_time = datetime.fromtimestamp(os.path.getctime(snapshot_path))
        if (now - created_time).days > retention_days:
            # Remove snapshot
            cmd = ["btrfs", "subvolume", "delete", snapshot_path]
            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            remove_elilo_entry(config, snapshot)
            print(f"Removed snapshot: {snapshot}")

# Adopt a snapshot as the new root subvolume
def adopt(config, snapshot):
    current_root_path = os.path.join(config['snapshot']['btrfs_mount_point'], config['snapshot']['root_subvolume'])
    snapshot_path = os.path.join(config['snapshot']['snapshot_dir'], snapshot)

    if not os.path.exists(snapshot_path):
        raise Exception(f"Snapshot '{snapshot}' not found.")

    # Step 1: Rename the current root subvolume to something else (e.g., @old_root)
    old_root_path = current_root_path + "_old"
    print(f"Renaming current root subvolume to {old_root_path}")
    cmd_rename_root = ["btrfs", "subvolume", "snapshot", current_root_path, old_root_path]
    subprocess.run(cmd_rename_root, check=True)

    # Step 2: Rename the selected snapshot to become the new root subvolume
    print(f"Renaming snapshot '{snapshot}' to become the new root subvolume.")
    cmd_set_new_root = ["btrfs", "subvolume", "snapshot", snapshot_path, current_root_path]
    subprocess.run(cmd_set_new_root, check=True)

    # Step 3: Update the elilo bootloader configuration to point to the new root subvolume
    print(f"Updating elilo.conf to use the new root subvolume '{snapshot}'.")
    add_elilo_entry(config, snapshot)

    print(f"Snapshot '{snapshot}' has been adopted as the new root subvolume.")

# Main logic for managing snapshots
def update(config):
    # Step 1: Get list of current snapshots and add them to elilo.conf if needed
    snapshots = list_snapshots(config)
    for snapshot in snapshots:
        add_elilo_entry(config, snapshot)

    # Step 2: Remove snapshots older than retention period
    remove_old_snapshots(config)

# Parse command-line arguments
def parse_arguments():
    parser = argparse.ArgumentParser(description='Manage BTRFS snapshots and elilo bootloader entries.')
    parser.add_argument('action', choices=['create', 'update', 'adopt'],
                        help='Specify the action to perform: create, manage, or adopt.')
    parser.add_argument('--snapshot', help='The snapshot to adopt as the new root subvolume.')
    return parser.parse_args()

def main():
    # Load configuration from /etc/manage_snapshots/config.toml
    config = read_config()

    # Parse command-line arguments
    args = parse_arguments()

    # Execute based on action
    if args.action == 'create':
        create(config)
    elif args.action == 'update':
        update(config)
    elif args.action == 'adopt':
        if not args.snapshot:
            print("Error: You must specify the snapshot to adopt using --snapshot.")
        else:
            adopt(config, args.snapshot)
    else:
        print("Invalid action")

if __name__ == "__main__":
    main()