import os
import subprocess
import toml
from datetime import datetime, timedelta
import shutil

# Read configuration from a TOML file
def read_config(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_snapshot(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):
    entry = f"""
image=vmlinuz
        label={snapshot}
        root={config['elilo']['root_partition']}
        append="rootflags=subvol={config['snapshot']['snapshot_dir']}/{snapshot} ro"
        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}")

# Function to wrap slackpkg upgrade
def slackpkg_upgrade(config):
    print("Preparing to run slackpkg upgrade...")
    create_snapshot(config)
    print("Running slackpkg upgrade...")
    subprocess.run(["/usr/sbin/slackpkg", "upgrade-all"])

# Function to wrap sbopkg -i
def sbopkg_install(config, package):
    print(f"Preparing to install package using sbopkg: {package}")
    create_snapshot(config)
    print(f"Running sbopkg -i {package}...")
    subprocess.run(["/usr/sbin/sbopkg", "-i", package])

def main():
    # Load configuration from config.toml
    config = read_config("config.toml")

    # 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)

if __name__ == "__main__":
    main()