/*
    * FsProbe.cpp
    * Filesystem probe registry
    * Copyright (c) 2026 Daniel Hammer
*/

#include "FsProbe.hpp"
#include <Drivers/Storage/Gpt.hpp>
#include <Drivers/Storage/BlockDevice.hpp>
#include <Terminal/Terminal.hpp>
#include <Libraries/Memory.hpp>

namespace Fs::FsProbe {

    static ProbeFn g_probes[MaxProbes] = {};
    static int g_probeCount = 0;
    static bool g_mounted[Drivers::Storage::Gpt::MaxPartitions] = {};
    static int g_driveForPart[Drivers::Storage::Gpt::MaxPartitions] = {};

    void Register(ProbeFn fn) {
        if (g_probeCount < MaxProbes && fn) {
            g_probes[g_probeCount++] = fn;
        }
    }

    static bool IsEfiPartition(const Drivers::Storage::Gpt::PartitionInfo* part) {
        auto& a = part->TypeGuid;
        auto& b = Drivers::Storage::Gpt::GUID_EFI_SYSTEM;
        return a.Data1 == b.Data1 && a.Data2 == b.Data2 && a.Data3 == b.Data3 &&
               memcmp(a.Data4, b.Data4, 8) == 0;
    }

    static int TryMountPartitionAtLowestDrive(int partIndex, int firstDrive, bool efiLog) {
        if (!Vfs::IsInitialized()) return 0;
        if (partIndex < 0 || partIndex >= Drivers::Storage::Gpt::MaxPartitions) return 0;
        if (g_mounted[partIndex]) return 0;

        auto* part = Drivers::Storage::Gpt::GetPartition(partIndex);
        if (!part) return 0;

        int driveNum = Vfs::FindLowestAvailableDrive(firstDrive);
        if (driveNum < 0) {
            Kt::KernelLogStream(Kt::WARNING, "FsProbe") << "No free drive slot for partition "
                << partIndex;
            return -1;
        }

        for (int p = 0; p < g_probeCount; p++) {
            Vfs::FsDriver* driver = g_probes[p](
                part->BlockDevIndex, part->StartLba, part->SectorCount);

            if (!driver) continue;

            if (Vfs::RegisterDrive(driveNum, driver) == 0) {
                g_mounted[partIndex] = true;
                g_driveForPart[partIndex] = driveNum;
                Kt::KernelLogStream(Kt::OK, "FsProbe")
                    << (efiLog ? "Mounted EFI partition " : "Mounted partition ")
                    << partIndex << " as drive " << driveNum;
                return 1;
            }

            return -1;
        }

        Kt::KernelLogStream(Kt::INFO, "FsProbe") << "No filesystem recognized on partition "
            << partIndex;
        return 0;
    }

    void MountPartitions(int firstDrive) {
        int partCount = Drivers::Storage::Gpt::GetPartitionCount();
        if (partCount == 0 || g_probeCount == 0) return;

        // First pass: mount non-EFI partitions so they get lower drive numbers.
        for (int i = 0; i < partCount; i++) {
            auto* part = Drivers::Storage::Gpt::GetPartition(i);
            if (!part || IsEfiPartition(part)) continue;
            TryMountPartitionAtLowestDrive(i, firstDrive, false);
        }

        // Second pass: mount EFI system partitions after all others.
        for (int i = 0; i < partCount; i++) {
            auto* part = Drivers::Storage::Gpt::GetPartition(i);
            if (!part || !IsEfiPartition(part)) continue;
            TryMountPartitionAtLowestDrive(i, firstDrive, true);
        }
    }

    int MountPartitionsForBlockDevice(int blockDevIndex, int firstDrive) {
        if (!Vfs::IsInitialized() || g_probeCount == 0) return 0;

        int partCount = Drivers::Storage::Gpt::GetPartitionCount();
        int mounted = 0;

        for (int i = 0; i < partCount; i++) {
            auto* part = Drivers::Storage::Gpt::GetPartition(i);
            if (!part || part->BlockDevIndex != blockDevIndex || IsEfiPartition(part)) continue;

            int result = TryMountPartitionAtLowestDrive(i, firstDrive, false);
            if (result > 0) mounted += result;
        }

        for (int i = 0; i < partCount; i++) {
            auto* part = Drivers::Storage::Gpt::GetPartition(i);
            if (!part || part->BlockDevIndex != blockDevIndex || !IsEfiPartition(part)) continue;

            int result = TryMountPartitionAtLowestDrive(i, firstDrive, true);
            if (result > 0) mounted += result;
        }

        return mounted;
    }

    int MountPartition(int partIndex, int driveNum) {
        if (driveNum < 0 || driveNum >= Vfs::MaxDrives) return -1;
        if (g_probeCount == 0) return -1;
        if (partIndex >= 0 && partIndex < Drivers::Storage::Gpt::MaxPartitions && g_mounted[partIndex]) return -1;
        if (Vfs::IsDriveRegistered(driveNum)) return -1;

        auto* part = Drivers::Storage::Gpt::GetPartition(partIndex);
        if (!part) return -1;

        for (int p = 0; p < g_probeCount; p++) {
            Vfs::FsDriver* driver = g_probes[p](
                part->BlockDevIndex, part->StartLba, part->SectorCount);

            if (driver) {
                if (Vfs::RegisterDrive(driveNum, driver) == 0) {
                    if (partIndex >= 0 && partIndex < Drivers::Storage::Gpt::MaxPartitions) {
                        g_mounted[partIndex] = true;
                        g_driveForPart[partIndex] = driveNum;
                    }
                    Kt::KernelLogStream(Kt::OK, "FsProbe") << "Mounted partition "
                        << partIndex << " as drive " << driveNum;
                    return 0;
                }
                return -1;
            }
        }

        Kt::KernelLogStream(Kt::WARNING, "FsProbe") << "No filesystem recognized on partition " << partIndex;
        return -1;
    }

    int GetBlockDeviceForDrive(int driveNum) {
        if (driveNum < 0) return -1;
        int partCount = Drivers::Storage::Gpt::GetPartitionCount();
        for (int i = 0; i < partCount && i < Drivers::Storage::Gpt::MaxPartitions; i++) {
            if (g_mounted[i] && g_driveForPart[i] == driveNum) {
                auto* part = Drivers::Storage::Gpt::GetPartition(i);
                if (part) return part->BlockDevIndex;
            }
        }
        return -1;
    }

    int UnmountPartitionsForBlockDevice(int blockDevIndex) {
        int partCount = Drivers::Storage::Gpt::GetPartitionCount();
        int unmounted = 0;
        int dst = 0;

        for (int src = 0; src < partCount; src++) {
            auto* part = Drivers::Storage::Gpt::GetPartition(src);
            bool remove = part && part->BlockDevIndex == blockDevIndex;

            if (remove) {
                if (src < Drivers::Storage::Gpt::MaxPartitions &&
                    g_mounted[src] && g_driveForPart[src] >= 0) {
                    if (Vfs::UnregisterDrive(g_driveForPart[src]) == 0) {
                        unmounted++;
                    }
                }
                if (src < Drivers::Storage::Gpt::MaxPartitions) {
                    g_mounted[src] = false;
                    g_driveForPart[src] = -1;
                }
                continue;
            }

            if (src < Drivers::Storage::Gpt::MaxPartitions &&
                dst < Drivers::Storage::Gpt::MaxPartitions) {
                if (dst != src) {
                    g_mounted[dst] = g_mounted[src];
                    g_driveForPart[dst] = g_driveForPart[src];
                }
                dst++;
            }
        }

        for (int i = dst; i < partCount && i < Drivers::Storage::Gpt::MaxPartitions; i++) {
            g_mounted[i] = false;
            g_driveForPart[i] = -1;
        }

        return unmounted;
    }

    int SyncAndUnmountAll() {
        // 1. Flush every block device's volatile write cache to media. FAT32 and
        //    ext2 are write-through, so file data already reached the device; the
        //    flush makes the device commit it to non-volatile storage.
        int devCount = Drivers::Storage::GetBlockDeviceCount();
        for (int i = 0; i < devCount; i++) {
            const auto* dev = Drivers::Storage::GetBlockDevice(i);
            if (dev && dev->Flush) {
                dev->Flush(dev->Ctx);
            }
        }

        // 2. Cleanly unmount every disk-backed VFS drive. The ramdisk has no
        //    backing block device (GetBlockDeviceForDrive < 0) and is left
        //    mounted so the rest of userspace keeps running until power-off.
        int drives[Vfs::MaxDrives];
        int n = Vfs::VfsDriveList(drives, Vfs::MaxDrives);
        int unmounted = 0;

        for (int i = 0; i < n; i++) {
            int driveNum = drives[i];
            if (GetBlockDeviceForDrive(driveNum) < 0) continue;

            if (Vfs::UnregisterDrive(driveNum) == 0) {
                unmounted++;
                for (int p = 0; p < Drivers::Storage::Gpt::MaxPartitions; p++) {
                    if (g_mounted[p] && g_driveForPart[p] == driveNum) {
                        g_mounted[p] = false;
                        g_driveForPart[p] = -1;
                    }
                }
            }
        }

        Kt::KernelLogStream(Kt::OK, "FsProbe") << "Synced and unmounted "
            << unmounted << " disk-backed volume(s) for power-off";
        return unmounted;
    }

};