import subprocess
import time

def run_command(cmd):
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Command failed: {cmd}\n{result.stderr}")
        return ""
    return result.stdout


def get_active_snapshot_keys():
    """Build a snapshot -> key map from `ctr snapshots list`."""
    output = run_command("ctr snapshots list")
    active_snapshots = []

    for line in output.splitlines():
        # line format is typically: SNAPSHOT-NAME  KIND  SIZE  INODES  CREATED AT  UPDATED AT  KEY
        parts = line.split()
        if len(parts) < 3:
            continue
        kind = parts[2]
        key = parts[0]
        if kind != "Active" or key == "overlayfs":
            continue
        else:
            active_snapshots.append(key)
    return active_snapshots


def get_snapshot_tree():
    """Build parent -> child relationship from `ctr snapshots tree`."""
    output = run_command("ctr snapshots tree")
    tree = {}
    current_root = None
    last_indent_pos = -1

    for line in output.splitlines():
        if not line.strip() or line.strip() == "overlayfs":
            continue

        indent_pos = line.find("\\_")
        snapshot_id = line.strip().replace("\\_", "").strip()

        if indent_pos == -1:
            # No \_, treat as a top-level root
            current_root = snapshot_id
            tree[current_root] = []
            last_indent_pos = -1
            continue

        if indent_pos <= last_indent_pos:
            # Start a new root chain
            current_root = snapshot_id
            tree[current_root] = []
        else:
            # Child of the current root
            if current_root:
                tree[current_root].append(snapshot_id)

        last_indent_pos = indent_pos

    return tree


def remove_snapshot(stale_snapshots):
    for snapshot in stale_snapshots:
        print(f"Removing snapshot {snapshot}")
        result = run_command(f"ctr snapshots rm {snapshot}")
        if result:
            print(result)

def get_stale_snapshots(tree, active_snapshots):
    stale_snapshots = []
    for parent, children in tree.items():
        if children[len(children) - 1] not in active_snapshots:
            stale_snapshots.append(children[len(children) - 1])
    return stale_snapshots

def main():
    run_command("systemctl start containerd")
    size = run_command("du -hs /var/lib/containerd/io.containerd.snapshotter.v1.overlayfs/")
    print("Snapshotter path size:\n", size)
    active_snapshots = get_active_snapshot_keys()
    print("Active snapshots:\n", active_snapshots)
    tree = get_snapshot_tree()
    for root, children in tree.items():
        print(f"{root}:")
        for child in children:
            print(f"  {child}")
    print("Snapshot tree:\n")
    stale_snapshots = get_stale_snapshots(tree, active_snapshots)
    if stale_snapshots:
        print("Stale snapshots:\n", stale_snapshots)
        remove_snapshot(stale_snapshots)
        time.sleep(10)
        tree = get_snapshot_tree()
        print("Snapshot tree after cleanup:\n")
        for root, children in tree.items():
            print(f"{root}:")
            for child in children:
                print(f"  {child}")
        size = run_command("du -hs /var/lib/containerd/io.containerd.snapshotter.v1.overlayfs/")
        print("Snapshotter path size after cleanup:\n", size)
    run_command("systemctl stop containerd")


if __name__ == "__main__":
    main()