using System.Management;
using System.Runtime.Versioning;
using System.Text.Json;
using DokanNet;
using DokanNet.Logging;

namespace DevDrive;

/// <summary>
/// Entry point for the virtual drive application.
/// Monitors a config drive via WMI events and creates symlinks based on index.json.
/// </summary>
[SupportedOSPlatform("windows")]
public static class Program
{
    private const string DefaultMountPoint = "A:\\";
    private const string DefaultConfigDrive = "C:\\PhysicalDrive\\";
    private const string ConfigFile = "index.json";
    private const string LogPrefix = "[DevDrive]";

    private static Dokan? _dokan;
    private static DokanInstance? _dokanInstance;
    private static readonly ManualResetEvent ShutdownEvent = new(false);
    private static bool _isMounted;
    private static string _configDrive = DefaultConfigDrive;
    private static string _configFilePath = "";
    private static string _mountPoint = DefaultMountPoint;
    private static FloppyConfig? _lastConfig;
    private static readonly object MountLock = new();
    private static bool _isConfigDriveLetter;

    /// <summary>
    /// Main entry point.
    /// </summary>
    /// <param name="args">Command line arguments. [0] = mount point, [1] = config drive.</param>
    public static void Main(string[] args)
    {
        _mountPoint = args.Length > 0 ? args[0] : DefaultMountPoint;
        _configDrive = args.Length > 1 ? args[1] : DefaultConfigDrive;

        if (!_mountPoint.EndsWith("\\"))
        {
            _mountPoint += "\\";
        }
        if (!_configDrive.EndsWith("\\"))
        {
            _configDrive += "\\";
        }

        _configFilePath = Path.Combine(_configDrive, ConfigFile);
        _isConfigDriveLetter = _configDrive.Length <= 3 && _configDrive[1] == ':';

        Console.WriteLine($"{LogPrefix} DevDrive");
        Console.WriteLine($"{LogPrefix} ==================");
        Console.WriteLine($"{LogPrefix} Watching for config drive: {_configDrive}");
        Console.WriteLine($"{LogPrefix} Config file: {_configFilePath}");
        Console.WriteLine($"{LogPrefix} Mount point: {_mountPoint}");
        Console.WriteLine();
        Console.WriteLine($"{LogPrefix} Press Ctrl+C to exit.");
        Console.WriteLine();

        Console.CancelKeyPress += (sender, eventArgs) =>
        {
            eventArgs.Cancel = true;
            Console.WriteLine($"{LogPrefix} Shutting down...");
            ShutdownEvent.Set();
        };

        try
        {
            _dokan = new Dokan(new ConsoleLogger(LogPrefix));
            
            if (_isConfigDriveLetter)
            {
                StartWmiWatchers();
            }
            else
            {
                StartFileSystemWatcher();
            }
            
            CheckAndMountIfReady();
            
            ShutdownEvent.WaitOne();
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Fatal error: {ex.Message}");
            Environment.Exit(1);
        }
        finally
        {
            UnmountFloppy();
            _dokan?.Dispose();
        }
    }

    private static void StartFileSystemWatcher()
    {
        StartVolumeWatchers();

        Console.WriteLine($"{LogPrefix} Volume watchers started for {_configDrive}");
        Console.WriteLine($"{LogPrefix} Watching for: {ConfigFile}");
    }

    private static void StartVolumeWatchers()
    {
        try
        {
            var insertQuery = new WqlEventQuery(
                "SELECT * FROM __InstanceCreationEvent WITHIN 1 WHERE TargetInstance ISA 'Win32_Volume'"
            );
            var insertWatcher = new ManagementEventWatcher(insertQuery);
            insertWatcher.EventArrived += (s, e) =>
            {
                Console.WriteLine($"{LogPrefix} Volume added, checking config...");
                Thread.Sleep(1000);
                CheckAndMountIfReady();
            };
            insertWatcher.Start();

            var removeQuery = new WqlEventQuery(
                "SELECT * FROM __InstanceDeletionEvent WITHIN 1 WHERE TargetInstance ISA 'Win32_Volume'"
            );
            var removeWatcher = new ManagementEventWatcher(removeQuery);
            removeWatcher.EventArrived += (s, e) =>
            {
                Console.WriteLine($"{LogPrefix} Volume removed");
                lock (MountLock)
                {
                    if (_isMounted && !IsConfigAccessible())
                    {
                        Console.WriteLine($"{LogPrefix} Config drive no longer accessible, unmounting...");
                        UnmountFloppy();
                        _lastConfig = null;
                    }
                }
            };
            removeWatcher.Start();
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Warning: Could not start WMI volume watchers: {ex.Message}");
        }
    }

    private static bool IsConfigAccessible()
    {
        try
        {
            if (!Directory.Exists(_configDrive))
            {
                return false;
            }

            var _ = Directory.GetFiles(_configDrive);
            return File.Exists(_configFilePath);
        }
        catch
        {
            return false;
        }
    }

    private static void StartWmiWatchers()
    {
        var driveLetter = _configDrive.TrimEnd('\\');

        var insertQuery = new WqlEventQuery(
            "SELECT * FROM __InstanceCreationEvent WITHIN 1 WHERE TargetInstance ISA 'Win32_LogicalDisk'"
        );
        var insertWatcher = new ManagementEventWatcher(insertQuery);
        insertWatcher.EventArrived += OnDriveInserted;
        insertWatcher.Start();

        var removeQuery = new WqlEventQuery(
            "SELECT * FROM __InstanceDeletionEvent WITHIN 1 WHERE TargetInstance ISA 'Win32_LogicalDisk'"
        );
        var removeWatcher = new ManagementEventWatcher(removeQuery);
        removeWatcher.EventArrived += OnDriveRemoved;
        removeWatcher.Start();

        Console.WriteLine($"{LogPrefix} WMI watchers started (event-based, no polling)");
        Console.WriteLine($"{LogPrefix} Watching for drive: {driveLetter}");
    }

    private static void OnDriveInserted(object sender, EventArrivedEventArgs e)
    {
        try
        {
            var targetInstance = (ManagementBaseObject)e.NewEvent["TargetInstance"];
            var deviceId = targetInstance["DeviceID"]?.ToString();

            if (string.IsNullOrEmpty(deviceId))
            {
                return;
            }

            var insertedDrive = deviceId + "\\";
            
            if (string.Equals(insertedDrive, _configDrive, StringComparison.OrdinalIgnoreCase))
            {
                Console.WriteLine();
                Console.WriteLine($"{LogPrefix} Drive {deviceId} inserted!");
                
                Thread.Sleep(500);
                CheckAndMountIfReady();
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Error handling drive insert: {ex.Message}");
        }
    }

    private static void OnDriveRemoved(object sender, EventArrivedEventArgs e)
    {
        try
        {
            var targetInstance = (ManagementBaseObject)e.NewEvent["TargetInstance"];
            var deviceId = targetInstance["DeviceID"]?.ToString();

            if (string.IsNullOrEmpty(deviceId))
            {
                return;
            }

            var removedDrive = deviceId + "\\";
            
            if (string.Equals(removedDrive, _configDrive, StringComparison.OrdinalIgnoreCase))
            {
                Console.WriteLine();
                Console.WriteLine($"{LogPrefix} Drive {deviceId} removed!");
                
                lock (MountLock)
                {
                    if (_isMounted)
                    {
                        Console.WriteLine($"{LogPrefix} Config drive disconnected, unmounting...");
                        UnmountFloppy();
                        _lastConfig = null;
                    }
                }
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Error handling drive removal: {ex.Message}");
        }
    }

    private static void CheckAndMountIfReady()
    {
        lock (MountLock)
        {
            var driveExists = Directory.Exists(_configDrive);
            var configExists = driveExists && File.Exists(_configFilePath);

            if (!driveExists)
            {
                Console.WriteLine($"{LogPrefix} Waiting for {_configDrive} to be connected...");
                return;
            }

            if (!configExists)
            {
                Console.WriteLine($"{LogPrefix} Drive present but {ConfigFile} not found");
                return;
            }

            var config = ReadConfig(_configFilePath);

            if (config == null || config.Paths.Count == 0)
            {
                Console.WriteLine($"{LogPrefix} Config empty or invalid");
                
                if (_isMounted)
                {
                    UnmountFloppy();
                    _lastConfig = null;
                }
                return;
            }

            var configChanged = !ConfigEquals(_lastConfig, config);

            if (!_isMounted || configChanged)
            {
                if (_isMounted && configChanged)
                {
                    Console.WriteLine($"{LogPrefix} Config changed, remounting...");
                    UnmountFloppy();
                }

                MountFloppy(config);
                _lastConfig = config;
            }
        }
    }

    private static FloppyConfig? ReadConfig(string configFilePath)
    {
        try
        {
            var json = File.ReadAllText(configFilePath);
            return JsonSerializer.Deserialize<FloppyConfig>(json);
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Error reading config: {ex.Message}");
            return null;
        }
    }

    private static bool ConfigEquals(FloppyConfig? a, FloppyConfig? b)
    {
        if (a == null || b == null)
        {
            return a == b;
        }

        if (a.Paths.Count != b.Paths.Count)
        {
            return false;
        }

        for (var i = 0; i < a.Paths.Count; i++)
        {
            if (a.Paths[i].Name != b.Paths[i].Name || a.Paths[i].Path != b.Paths[i].Path)
            {
                return false;
            }
        }

        return true;
    }

    private static void MountFloppy(FloppyConfig config)
    {
        if (_dokan == null)
        {
            return;
        }

        try
        {
            var symlinks = new Dictionary<string, string>();
            Console.WriteLine();
            Console.WriteLine($"{LogPrefix} Creating symlinks:");

            foreach (var entry in config.Paths)
            {
                var symlinkPath = "\\" + entry.Name;
                symlinks[symlinkPath] = entry.Path;
                Console.WriteLine($"{LogPrefix}   {entry.Name} -> {entry.Path}");
            }

            Console.WriteLine();

            var totalBytes = (long)(config.DriveTotalMB * 1024 * 1024);
            var freeBytes = (long)(config.DriveFreeMB * 1024 * 1024);
            Console.WriteLine($"{LogPrefix} Drive size: {config.DriveTotalMB} MB total, {config.DriveFreeMB} MB free");
            var floppyFs = new FloppyFileSystem(symlinks, config.DriveName, totalBytes, freeBytes);

            var dokanBuilder = new DokanInstanceBuilder(_dokan)
                .ConfigureOptions(options =>
                {
                    options.Options = DokanOptions.MountManager | DokanOptions.CurrentSession;
                    options.MountPoint = _mountPoint;
                    options.SingleThread = false;
                    options.AllocationUnitSize = FloppyStorage.AllocationUnitSize;
                    options.SectorSize = FloppyStorage.SectorSize;
                });

            _dokanInstance = dokanBuilder.Build(floppyFs);
            _isMounted = true;

            Console.WriteLine($"{LogPrefix} Mounted at {_mountPoint}");

            var iconPath = string.IsNullOrWhiteSpace(config.DriveIcon)
                ? "%SystemRoot%\\System32\\shell32.dll,6"
                : config.DriveIcon;
            DriveIconHelper.SetDriveIcon(_mountPoint, iconPath, config.DriveName);
        }
        catch (DokanException ex)
        {
            Console.WriteLine($"{LogPrefix} Mount error: {ex.Message}");
            _isMounted = false;
        }
    }

    private static void UnmountFloppy()
    {
        if (!_isMounted)
        {
            return;
        }

        try
        {
            DriveIconHelper.RemoveDriveIcon();
            _dokan?.RemoveMountPoint(_mountPoint);
            _dokanInstance?.Dispose();
            _dokanInstance = null;
            _isMounted = false;
            Console.WriteLine($"{LogPrefix} Unmounted");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Unmount error: {ex.Message}");
        }
    }
}