using System.Runtime.Versioning;
using Microsoft.Win32;

namespace DevDrive;

/// <summary>
/// Sets custom drive icons via the Windows registry.
/// </summary>
[SupportedOSPlatform("windows")]
public static class DriveIconHelper
{
    private const string LogPrefix = "[DriveIcon]";
    private const string DriveIconsKey = @"Software\Classes\Applications\Explorer.exe\Drives";
    
    private static string? _currentDriveLetter;

    public static bool SetDriveIcon(string driveLetter, string iconPath, string? driveLabel = null)
    {
        try
        {
            var letter = driveLetter.TrimEnd('\\', ':').ToUpperInvariant();
            if (letter.Length != 1 || letter[0] < 'A' || letter[0] > 'Z')
            {
                Console.WriteLine($"{LogPrefix} Invalid drive letter: {driveLetter}");
                return false;
            }

            var driveKeyPath = @$"{DriveIconsKey}\{letter}\DefaultIcon";
            
            using var key = Registry.CurrentUser.CreateSubKey(driveKeyPath);
            if (key == null)
            {
                Console.WriteLine($"{LogPrefix} Failed to create registry key");
                return false;
            }

            key.SetValue("", iconPath);
            _currentDriveLetter = letter;

            if (!string.IsNullOrEmpty(driveLabel))
            {
                var labelKeyPath = @$"{DriveIconsKey}\{letter}\DefaultLabel";
                using var labelKey = Registry.CurrentUser.CreateSubKey(labelKeyPath);
                labelKey?.SetValue("", driveLabel);
            }

            RefreshShell();

            Console.WriteLine($"{LogPrefix} Set icon for drive {letter}: to {iconPath}");
            return true;
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Error setting icon: {ex.Message}");
            return false;
        }
    }

    public static void RemoveDriveIcon()
    {
        if (_currentDriveLetter == null)
        {
            return;
        }

        try
        {
            var driveKeyPath = @$"{DriveIconsKey}\{_currentDriveLetter}";
            Registry.CurrentUser.DeleteSubKeyTree(driveKeyPath, throwOnMissingSubKey: false);
            
            Console.WriteLine($"{LogPrefix} Removed icon for drive {_currentDriveLetter}:");
            _currentDriveLetter = null;

            RefreshShell();
        }
        catch (Exception ex)
        {
            Console.WriteLine($"{LogPrefix} Error removing icon: {ex.Message}");
        }
    }

    private static void RefreshShell()
    {
        try
        {
            NativeMethods.SHChangeNotify(
                NativeMethods.SHCNE_ASSOCCHANGED,
                NativeMethods.SHCNF_IDLIST,
                IntPtr.Zero,
                IntPtr.Zero);
        }
        catch
        {
        }
    }

    private static class NativeMethods
    {
        public const int SHCNE_ASSOCCHANGED = 0x08000000;
        public const int SHCNF_IDLIST = 0x0000;

        [System.Runtime.InteropServices.DllImport("shell32.dll")]
        public static extern void SHChangeNotify(int wEventId, int uFlags, IntPtr dwItem1, IntPtr dwItem2);
    }
}