package org.limine.snapper;

import org.limine.snapper.formats.json1.Manifest;
import org.limine.snapper.formats.json1.SnapshotEntry;
import org.limine.snapper.objects.Config;
import org.limine.snapper.objects.FileDetail;
import org.limine.snapper.objects.PartitionUsage;
import org.limine.snapper.objects.RestoreDetail;
import org.limine.snapper.processes.*;

import java.util.*;
import java.util.regex.Pattern;

public class Main {

    static void main(String[] args) {
        if (args.length == 0) {
            update();
            return;
        }

        String cmd = "--update";
        LinkedHashSet<String> params = new LinkedHashSet<>();

        List<String> list = Arrays.asList(args);
        Iterator<String> it = list.iterator();

        while (it.hasNext()) {
            String a = it.next();
            switch (a) {
                case "--no-force-save":
                    Config.IS_SAVE_FORCED = false;
                    break;
                case "--info":
                case "--list":
                case "--restore":
                    cmd = a;
                    break;
                case "--restore-kernels":
                    cmd = a;
                    if (it.hasNext()) {
                        String next = it.next();
                        if (!next.startsWith("--")) {
                            params.add(next);
                        }
                    }
                    break;
                case "--remove":
                    cmd = a;
                    while (it.hasNext()) {
                        String next = it.next();
                        if (!next.startsWith("--")) {
                            params.add(next);
                        } else {
                            break;
                        }
                    }
                    break;
                default:
                    // ignore unknown options
                    break;
            }
        }
        handleCommand(cmd, params);
    }

    private static void handleCommand(String cmd, LinkedHashSet<String> params) {
        switch (cmd) {
            case "--remove":
                if (!params.isEmpty()) {
                    removeEntries(params);
                } else {
                    Utility.warnMessage("Please give IDs like 1 3 5..7");
                }
                break;
            case "--restore-kernels":
                if (!params.isEmpty()) {
                    restoreKernels(params.getFirst());
                } else {
                    Utility.exitWithError("--restore-kernels requires snapshot ID.");
                }
                break;
            case "--restore":
                restore();
                break;
            case "--info":
                info();
                break;
            case "--list":
                list();
                break;
            default:
                update();
                break;
        }
    }

    private static void info() {
        // Load keywords and values from multiple config files
        Config config = new ConfigReader().readConfig();

        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);
        HistoryReader.displayInfo(manifest, config);
    }

    private static void list() {
        // Load keywords and values from multiple config files
        Config config = new ConfigReader().readConfig();

        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);
        HistoryReader.displaySnapshotList(manifest);
    }

    private static void removeEntries(HashSet<String> params) {
        LinkedHashSet<Integer> ids = new LinkedHashSet<>();
        final int maxIds = 1000;
        for (String param : params) {
            // Check global max limit before anything
            if (ids.size() >= maxIds) {
                Utility.warnMessage("Too many IDs (max " + maxIds + ")");
                return;
            }
            try {
                if (param.contains("..") || param.contains("-")) {
                    String sep = param.contains("..") ? ".." : "-";
                    String[] parts = param.split(Pattern.quote(sep));
                    if (parts.length == 2) {
                        int start = Integer.parseInt(parts[0]);
                        int end = Integer.parseInt(parts[1]);
                        int min = Math.min(start, end);
                        int max = Math.max(start, end);
                        int rangeSize = max - min + 1;

                        // Prevent exploding range
                        if (rangeSize > maxIds) {
                            Utility.warnMessage("Range too large (max " + maxIds + " IDs allowed)");
                            return;
                        }
                        for (int i = min; i <= max; i++) {
                            ids.add(i);
                        }
                    }
                } else {
                    ids.add(Integer.parseInt(param));
                }
            } catch (NumberFormatException e) {
                //Utility.warnMessage("Invalid number in range: " + param);
            }
        }

        ConfigReader configReader = new ConfigReader();
        configReader.loadQuietMode();
        Config config = configReader.readConfig();
        LimineReader limine = new LimineReader(config);
        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);
        // If the manifest version is old, it will be upgraded automatically
        boolean isManifestUpgraded = FormatConverter.upgradeManifest(manifest);

        int snapshotCountBefore = manifest.getSnapshotEntries().size();
        SnapshotManager manager = new SnapshotManager(config);
        HashSet<FileDetail> deletedFileDetails = manager.remove(manifest, ids);
        int snapshotCountAfter = manifest.getSnapshotEntries().size();

        if (deletedFileDetails.isEmpty()) {
            if (snapshotCountBefore == snapshotCountAfter) {
                Utility.infoMessage("No snapshot boot entries found");
            } else {
                Utility.infoMessage("No kernel files removed (still used by other snapshots).");
            }
        } else {
            long bytes = 0L;
            for (FileDetail file : deletedFileDetails) {
                bytes += file.bytes();
            }
            String fileInfo = "1 kernel file removed";
            if (deletedFileDetails.size() > 1) {
                fileInfo = deletedFileDetails.size() + " kernel files removed";
            }
            Utility.infoMessage("Freed " + FileDetail.toMiB(bytes) + " on the boot partition (" + fileInfo + ")");
        }

        if (isManifestUpgraded || manager.checkUpdate()) {
            HistoryWriter.backup(manifestPath);
            HistoryWriter.save(manifest, manifestPath);
            LimineWriter writer = new LimineWriter(config, limine);
            writer.backup();
            writer.setSnapshots(manifest);
            writer.save();
        }
    }

    private static void update() {
        if (!Utility.areYouInBtrfs()) {
            Utility.warnMessage("You are not in Btrfs, the tool does not update anything!");
            return;
        }

        // Load keywords and values from multiple config files
        ConfigReader configReader = new ConfigReader();
        configReader.loadQuietMode();
        Config config = configReader.readConfig();
        configReader.checkMountPaths(config);
        // Run SnapperReader to read a Snapper list
        configReader.readSnapperConfig();

        // If you boot into any snapshot, do not update anything!
        if (Utility.areYouInSnapshot() && !Config.USE_OPENSUSE_LAYOUT) {
            Utility.warnMessage("You are in the snapshot, the tool does not update anything!");
            return;
        }

        SnapperReader snapper = new SnapperReader();

        // Run LimineReader to read the Limine config and convert all entries to a new tree node.
        LimineReader limine = new LimineReader(config);

        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);
        // If the manifest is old, convert it to a new version.
        boolean isManifestUpgraded = FormatConverter.upgradeManifest(manifest);

        // Sync the Limine snapshots with the Snapper list
        SnapshotManager manager = new SnapshotManager(config);
        manager.sync(manifest, snapper);

        boolean updateLimineConfig = manager.checkUpdate();

        if (updateLimineConfig) {

            // Check if your target OS name is found in limine.conf
            if (limine.isTargetOsNotFound()) {
                String errMessage = "Please set TARGET_OS_NAME in " + Config.DEFAULT_CONFIG_FILE_PATH;
                if (Config.TARGET_OS_NAME != null) {
                    errMessage = "Target OS name '" + Config.TARGET_OS_NAME + "' was not found, or TARGET_OS_NAME is invalid in " + Config.DEFAULT_CONFIG_FILE_PATH;
                }
                Utility.errorMessage(errMessage);
                return;
            }

            // Check if your OS entry has at least one kernel node.
            if (limine.getTargetOsNode().nodes.isEmpty() || (limine.getTargetOsNode().nodes.size() < 2 && limine.isSnapshotsNodeInsideOS())) {
                String errMessage = "Your OS entry has no kernel in " + Utility.buildPath(config.espPath(), Config.LIMINE_CONFIG_FILE) + ". Please add at least a kernel entry '//<kernel name>'";
                Utility.errorMessage(errMessage);
                return;
            }

            if (manager.isAllowedToAddSnapshot()) {
                SnapshotEntry newSnapshotEntry = limine.createCurrentLimineEntry(snapper.getLastSnapperEntry());
                // Check the size of your boot partition if the usage limit is exceeded.
                PartitionUsage usage = manager.predictSpaceUsage(newSnapshotEntry);
                if (usage.exceedsLimit()) {
                    // Auto-prune older snapshot boot entries according to LIMIT_USAGE_PERCENT.
                    usage = manager.pruneSnapshotEntries(manifest, usage);
                }
                if (usage.exceedsLimit()) {
                    if (usage.percent(usage.usedBytes()) >= Config.LIMIT_USAGE_PERCENT) {
                        Utility.errorMessage("Stopping snapshot boot entry creation:", "Boot partition usage "
                                + usage.usedPercent() + " already exceeds the " + Config.LIMIT_USAGE_PERCENT + "% limit.");
                    } else if (usage.increaseAtLeastMiB(100)) {
                        Utility.errorMessage("Stopping snapshot boot entry creation:", "Boot partition usage will increase by " + usage.deltaMiB() + " from "
                                + usage.usedPercent() + " to " + usage.newPercent() + " and will exceed the " + Config.LIMIT_USAGE_PERCENT + "% limit.");
                    } else {
                        Utility.errorMessage("Stopping snapshot boot entry creation:", "Boot partition usage from "
                                + usage.usedPercent() + " to " + usage.newPercent() + " will exceed the " + Config.LIMIT_USAGE_PERCENT + "% limit.");
                    }
                } else {
                    if (Config.SNAPSHOT_WRITABLE) {
                        String makeSnapshotWritable = "btrfs property set '/.snapshots/" + newSnapshotEntry.snapperEntry().snapshotID() + "/snapshot' ro " + false;
                        Utility.runCommand(makeSnapshotWritable, false, false);
                    }
                    // The manifest adds the new snapshot.
                    manager.historyAddSnapshot(manifest, newSnapshotEntry);
                }
            }
        }

        if (updateLimineConfig || isManifestUpgraded || Config.FS_UUID != null && !Config.FS_UUID.equals(manifest.getUuid())) {
            // Overwrite /<esp>/<machine-id>/limine_history/snapshots.json
            HistoryWriter.backup(manifestPath);
            HistoryWriter.save(manifest, manifestPath);
        }

        if (Config.IS_SAVE_FORCED || updateLimineConfig) {
            // Save...
            LimineWriter writer = new LimineWriter(config, limine);
            // Backup /<esp>/limine.conf
            writer.backup();
            // Overwrite new snapshots and the current config in /<esp>/limine.conf
            writer.setSnapshots(manifest);
            writer.save();
        }
    }

    private static void restoreKernels(String snapshotID) {
        Config.ENABLE_TMP_CONFIG = Utility.areYouInBtrfs() && Utility.areYouInSnapshot();
        // Load keywords and values from multiple config files
        Config config = new ConfigReader().readConfig();

        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);

        // Check if your target OS name is found in limine.conf
        LimineReader limine = new LimineReader(config, false);
        if (limine.isTargetOsNotFound()) {
            String errMessage = "Please set TARGET_OS_NAME in " + Config.DEFAULT_CONFIG_FILE_PATH;
            if (Config.TARGET_OS_NAME != null) {
                errMessage = "Target OS name '" + Config.TARGET_OS_NAME + "' was not found, or TARGET_OS_NAME is invalid in " + Config.DEFAULT_CONFIG_FILE_PATH;
            }
            Utility.exitWithError(errMessage, false);
            return;
        }

        Config.IS_TIME_FOR_RESTORE = true;

        // Select snapshot ID for the restore.
        SnapshotManager manager = new SnapshotManager(config);

        // Check if a restore detail is valid
        RestoreDetail restoreDetail = manager.preRestoreKernels(manifest, limine, snapshotID);
        if (restoreDetail == null || !restoreDetail.isValid()) {
            Config.IS_TIME_FOR_RESTORE = false;
            Utility.exitWithError("The restore is invalid.", false);
            return;
        }

        // It's time to start restoring a snapshot.
        // Copy the selected snapshot entry from snapshots.json and write it to the target OS node.
        SnapshotEntry restoredEntry = manager.restore(manifest, limine, restoreDetail, true);
        if (restoredEntry == null) {
            Config.IS_TIME_FOR_RESTORE = false;
            Utility.exitWithError("The restore failed.", false);
        } else {
            HistoryWriter.backup(manifestPath);
            // Overwrite /<esp>/<machine-id>/limine_history/snapshots.json
            HistoryWriter.save(manifest, manifestPath);

            LimineWriter writer = new LimineWriter(config, limine);
            // Backup /<esp>/limine.conf
            writer.backup();
            // Overwrite the updated snapshots-list and the current configs in /<esp>/limine.conf
            writer.setTargetEntry(restoredEntry);
            writer.setSnapshots(manifest);
            writer.save();
        }
    }

    private static void restore() {
        Config.IS_TIME_FOR_RESTORE = true;
        Config.ENABLE_TMP_CONFIG = Utility.areYouInBtrfs() && Utility.areYouInSnapshot();
        // Load keywords and values from multiple config files
        Config config = new ConfigReader().readConfig();

        // Read /<esp>/<machine-id>/limine_history/snapshots.json
        String manifestPath = Utility.buildPath(config.espPath(), config.machineID(), Config.HISTORY_FOLDER, Config.MANIFEST_FILE);
        Manifest manifest = HistoryReader.load(manifestPath);

        // Check if your target OS name is found in limine.conf
        LimineReader limine = new LimineReader(config, true);
        if (limine.isTargetOsNotFound()) {
            String errMessage = "Please set TARGET_OS_NAME in " + Config.DEFAULT_CONFIG_FILE_PATH;
            if (Config.TARGET_OS_NAME != null) {
                errMessage = "Target OS name '" + Config.TARGET_OS_NAME + "' was not found, or TARGET_OS_NAME is invalid in " + Config.DEFAULT_CONFIG_FILE_PATH;
            }
            Utility.errorMessage(errMessage);
            return;
        }

        // Select snapshot ID for the restore.
        SnapshotManager manager = new SnapshotManager(config);

        // Check if a restore detail is valid
        RestoreDetail restoreDetail = manager.preRestore(manifest, limine);
        if (restoreDetail == null || !restoreDetail.isValid()) {
            Utility.exitWithError("The restore is invalid.", false);
            return;
        }

        if (Config.USE_OPENSUSE_LAYOUT) {
            Utility.deleteFile(Config.LOCK_RESTORE_FILE, false, false);
            String snapperRollbackCmdline = "snapper rollback " + restoreDetail.snapshotEntry().snapperEntry().snapshotID();
            if (!Utility.runCommand(snapperRollbackCmdline, true, true)) {
                Utility.exitWithError("snapper rollback failed.", true);
                return;
            }
        } else {
            // It's time to start restoring a snapshot.
            // Copy the selected snapshot entry from snapshots.json and write it to the target OS node.
            SnapshotEntry restoredEntry = manager.restore(manifest, limine, restoreDetail, false);

            if (restoredEntry == null) {
                Utility.exitWithError("The restore failed.", true);
                return;
            } else {
                HistoryWriter.backup(manifestPath);
                // Overwrite /<esp>/<machine-id>/limine_history/snapshots.json
                HistoryWriter.save(manifest, manifestPath);

                LimineWriter writer = new LimineWriter(config, limine);
                // Backup /<esp>/limine.conf
                writer.backup();
                // Overwrite the updated snapshots-list and the current configs in /<esp>/limine.conf
                writer.setTargetEntry(restoredEntry);
                writer.setSnapshots(manifest);
                writer.save();
            }
        }

        if (Utility.isCommandPresent("reboot")) {
            boolean needRestart = Utility.askUser("""
                    \s
                     The restore is complete. Please reboot now:
                     Type [y]es to reboot, [n]o or [c]ancel to abort rebooting.
                    \s""", "yes", "no", "cancel");
            if (needRestart) {
                Utility.runCommand("reboot", true, false);
            }
        } else {
            Utility.askUser("""
                    \s
                     The restore is complete. Please reboot now.
                     Type [y]es to close.
                    \s""", "yes", "no", "cancel");
        }
    }
}





