using System.ComponentModel;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Security.Principal;
using System.Text;
using Microsoft.Win32;

namespace ONE.CleanupAdmin;

internal static class Program
{
    private const string DefaultTargetDir = @"C:\Program Files (x86)\Internet Download Manager";
    private const string ChromeExtensionId = "ngpampappnmepgilojfohadhhmbhlaek";
    private const string RunOnceName = "ONE Clean IDM Leftovers";
    private const int MoveFileDelayUntilReboot = 4;

    private static readonly string[] ProcessNames =
    [
        "IDMan.exe",
        "IDMGrHlp.exe",
        "IDMIntegrator64.exe",
        "IDMMsgHost.exe",
        "explorer.exe"
    ];

    private static readonly RegistryTarget[] RegistryTargets =
    [
        new(RegistryHive.LocalMachine, RegistryView.Registry64, @"Software\Google\Chrome\NativeMessagingHosts\com.tonec.idm"),
        new(RegistryHive.LocalMachine, RegistryView.Registry64, @"Software\WOW6432Node\Google\Chrome\NativeMessagingHosts\com.tonec.idm"),
        new(RegistryHive.CurrentUser, RegistryView.Registry64, @"Software\Google\Chrome\NativeMessagingHosts\com.tonec.idm"),
        new(RegistryHive.CurrentUser, RegistryView.Registry64, @"Software\DownloadManager")
    ];

    public static int Main(string[] args)
    {
        Console.OutputEncoding = Encoding.UTF8;
        var options = Options.Parse(args);
        if (options.ShowHelp)
        {
            Options.PrintHelp();
            return 0;
        }

        using var log = LogWriter.Create(options.LogPath);
        try
        {
            var isAdmin = IsAdministrator();
            WriteHeader(options, isAdmin, log);

            if (!isAdmin && !options.DryRun)
            {
                if (options.Elevate && !options.DryRun)
                {
                    log.WriteLine("Starting elevated helper.");
                    return StartElevated(args);
                }

                log.WriteLine("ERROR: administrator rights are required. Use --dry-run or run elevated.");
                return 2;
            }

            var runner = new CleanupRunner(options, log);
            return runner.Run();
        }
        catch (Exception ex)
        {
            log.WriteLine("ERROR: " + ex.Message);
            log.WriteDetail("DETAIL: " + ex);
            WriteRecoveryGuidance(log);
            return 1;
        }
        finally
        {
            if (options.Pause)
            {
                log.WriteLine("Press any key to exit...");
                Console.ReadKey(intercept: true);
            }
        }
    }

    private static void WriteHeader(Options options, bool isAdmin, LogWriter log)
    {
        log.WriteLine("ONE CleanupAdmin");
        log.WriteLine("Mode: " + (options.DryRun ? "dry-run" : "execute"));
        log.WriteLine("Admin: " + (isAdmin ? "yes" : "no"));
        log.WriteLine("Target: " + options.TargetDir);
        log.WriteLine("Log: " + (log.FilePath ?? "disabled"));
        log.WriteLine();
    }

    private static void WriteRecoveryGuidance(LogWriter log)
    {
        log.WriteLine("Recovery guidance:");
        log.WriteLine("  1. If Explorer is closed, run explorer.exe from Task Manager or a command prompt.");
        log.WriteLine("  2. Reboot once if files were scheduled for deletion, then run --dry-run to inspect remaining targets.");
        log.WriteLine("  3. If registry/native messaging entries must be restored, reinstall the owning app or restore them from backup.");
        log.WriteLine("  4. Review the log before retrying elevated.");
    }

    private static int StartElevated(string[] args)
    {
        var exePath = Environment.ProcessPath ?? throw new InvalidOperationException("Cannot resolve helper path.");
        var start = new ProcessStartInfo
        {
            FileName = exePath,
            UseShellExecute = true,
            Verb = "runas",
            Arguments = string.Join(" ", args.Where(a => !a.Equals("--elevate", StringComparison.OrdinalIgnoreCase)).Select(QuoteArg))
        };
        Process.Start(start);
        return 0;
    }

    private static string QuoteArg(string value)
    {
        if (value.Length == 0) return "\"\"";
        if (!value.Any(char.IsWhiteSpace) && !value.Contains('"')) return value;
        return "\"" + value.Replace("\\", "\\\\").Replace("\"", "\\\"") + "\"";
    }

    private static bool IsAdministrator()
    {
        using var identity = WindowsIdentity.GetCurrent();
        var principal = new WindowsPrincipal(identity);
        return principal.IsInRole(WindowsBuiltInRole.Administrator);
    }

    [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
    private static extern bool MoveFileEx(string lpExistingFileName, string? lpNewFileName, int dwFlags);

    private sealed class CleanupRunner(Options options, LogWriter log)
    {
        public int Run()
        {
            ValidateTargetPath(options.TargetDir);
            ShowPlan();
            log.WriteLine();

            KillProcesses();
            DeleteRegistryTargets();
            UnregisterShellExtension();
            PrepareAndDeleteTargetDir();
            CleanPendingFileRenameOperations();
            DeleteChromeExtensionFolders();
            RegisterRunOnce();
            StartExplorer();

            log.WriteLine();
            log.WriteLine("IDM 관리자 권한 잔여물 정리 완료.");
            return 0;
        }

        private void ShowPlan()
        {
            log.WriteLine("Deletion targets:");
            log.WriteLine("  Processes:");
            foreach (var name in ProcessNames) log.WriteLine("    " + name);
            log.WriteLine("  Registry:");
            foreach (var target in RegistryTargets) log.WriteLine("    " + target);
            log.WriteLine("    HKLM\\SYSTEM\\CurrentControlSet\\Control\\Session Manager value PendingFileRenameOperations entries matching IDM patterns");
            log.WriteLine("  Files:");
            log.WriteLine("    " + options.TargetDir);
            log.WriteLine("    Chrome profile Extensions\\" + ChromeExtensionId + " under " + options.ChromeUserDataDir);
            log.WriteLine("  RunOnce:");
            log.WriteLine("    HKCU\\Software\\Microsoft\\Windows\\CurrentVersion\\RunOnce\\" + RunOnceName);
            log.WriteLine("  Explorer:");
            log.WriteLine("    restart explorer.exe");
        }

        private void KillProcesses()
        {
            foreach (var processName in ProcessNames)
            {
                RunProcess("taskkill.exe", ["/IM", processName, "/F"], "Kill process " + processName, ignoreExitCode: true);
            }
        }

        private void DeleteRegistryTargets()
        {
            foreach (var target in RegistryTargets)
            {
                RunRegistryAction("Delete registry " + target, () =>
                {
                    using var hive = RegistryKey.OpenBaseKey(target.Hive, target.View);
                    hive.DeleteSubKeyTree(target.SubKey, throwOnMissingSubKey: false);
                });
            }
        }

        private void UnregisterShellExtension()
        {
            var dllPath = Path.Combine(options.TargetDir, "IDMShellExt64.dll");
            if (!File.Exists(dllPath))
            {
                Log("Missing shell extension: " + dllPath);
                return;
            }

            RunProcess("regsvr32.exe", ["/s", "/u", dllPath], "Unregister " + dllPath, ignoreExitCode: false);
        }

        private void PrepareAndDeleteTargetDir()
        {
            if (!Directory.Exists(options.TargetDir))
            {
                Log("Missing target folder: " + options.TargetDir);
                return;
            }

            RunProcess("takeown.exe", ["/F", options.TargetDir, "/R", "/D", "Y"], "Take ownership " + options.TargetDir, ignoreExitCode: true);
            RunProcess("icacls.exe", [options.TargetDir, "/grant", "Administrators:F", "/T", "/C"], "Grant Administrators:F " + options.TargetDir, ignoreExitCode: true);
            NormalizeAttributes(options.TargetDir);
            RunFileAction("Delete target folder " + options.TargetDir, () => Directory.Delete(options.TargetDir, recursive: true));

            if (Directory.Exists(options.TargetDir))
            {
                ScheduleDeleteOnReboot(options.TargetDir);
            }
        }

        private void NormalizeAttributes(string root)
        {
            RunFileAction("Normalize attributes under " + root, () =>
            {
                foreach (var path in Directory.EnumerateFileSystemEntries(root, "*", SearchOption.AllDirectories))
                {
                    try
                    {
                        File.SetAttributes(path, FileAttributes.Normal);
                    }
                    catch (Exception ex) when (ex is UnauthorizedAccessException or IOException)
                    {
                        Log("Attribute skip: " + path + " (" + ex.Message + ")");
                    }
                }
            });
        }

        private void ScheduleDeleteOnReboot(string root)
        {
            RunFileAction("Schedule delete on reboot " + root, () =>
            {
                foreach (var path in Directory.EnumerateFileSystemEntries(root, "*", SearchOption.AllDirectories).OrderByDescending(p => p, StringComparer.OrdinalIgnoreCase))
                {
                    MoveFileExOrThrow(path);
                }

                MoveFileExOrThrow(root);
            });
        }

        private static void MoveFileExOrThrow(string path)
        {
            if (!MoveFileEx(path, null, MoveFileDelayUntilReboot))
            {
                throw new Win32Exception(Marshal.GetLastWin32Error(), "MoveFileEx failed for " + path);
            }
        }

        private void CleanPendingFileRenameOperations()
        {
            RunRegistryAction("Clean IDM PendingFileRenameOperations entries", () =>
            {
                using var hive = RegistryKey.OpenBaseKey(RegistryHive.LocalMachine, RegistryView.Registry64);
                using var key = hive.OpenSubKey(@"SYSTEM\CurrentControlSet\Control\Session Manager", writable: true);
                var value = key?.GetValue("PendingFileRenameOperations") as string[];
                if (key == null || value == null)
                {
                    Log("Missing PendingFileRenameOperations");
                    return;
                }

                var keep = new List<string>();
                for (var i = 0; i < value.Length; i += 2)
                {
                    var src = value[i];
                    var dst = i + 1 < value.Length ? value[i + 1] : "";
                    if (IsIdmPendingEntry(src) || IsIdmPendingEntry(dst))
                    {
                        Log("Remove pending pair: " + src + " -> " + dst);
                        continue;
                    }

                    keep.Add(src);
                    if (i + 1 < value.Length) keep.Add(dst);
                }

                if (keep.Count == 0)
                {
                    key.DeleteValue("PendingFileRenameOperations", throwOnMissingValue: false);
                }
                else
                {
                    key.SetValue("PendingFileRenameOperations", keep.ToArray(), RegistryValueKind.MultiString);
                }
            });
        }

        private static bool IsIdmPendingEntry(string value)
        {
            return value.Contains("Internet Download Manager", StringComparison.OrdinalIgnoreCase)
                || value.Contains(@"\IDM\", StringComparison.OrdinalIgnoreCase)
                || value.Contains("IDMan", StringComparison.OrdinalIgnoreCase)
                || value.Contains("tonec", StringComparison.OrdinalIgnoreCase)
                || value.Contains("DownloadManager", StringComparison.OrdinalIgnoreCase);
        }

        private void DeleteChromeExtensionFolders()
        {
            if (!Directory.Exists(options.ChromeUserDataDir))
            {
                Log("Missing Chrome user data: " + options.ChromeUserDataDir);
                return;
            }

            var found = false;
            foreach (var profileDir in Directory.EnumerateDirectories(options.ChromeUserDataDir))
            {
                var extensionDir = Path.Combine(profileDir, "Extensions", ChromeExtensionId);
                if (Directory.Exists(extensionDir))
                {
                    found = true;
                    RunFileAction("Delete Chrome extension folder " + extensionDir, () => Directory.Delete(extensionDir, recursive: true));
                }
            }

            if (!found)
            {
                Log("No Chrome extension folders found for " + ChromeExtensionId);
            }
        }

        private void RegisterRunOnce()
        {
            var command = "powershell.exe -NoProfile -ExecutionPolicy Bypass -Command \"Remove-Item -LiteralPath ''" + options.TargetDir.Replace("'", "''") + "'' -Recurse -Force -ErrorAction SilentlyContinue\"";
            RunRegistryAction("Register RunOnce " + RunOnceName + " = " + command, () =>
            {
                using var key = Registry.CurrentUser.CreateSubKey(@"Software\Microsoft\Windows\CurrentVersion\RunOnce", writable: true);
                key.SetValue(RunOnceName, command, RegistryValueKind.String);
            });
        }

        private void StartExplorer()
        {
            RunProcess("explorer.exe", [], "Start explorer.exe", ignoreExitCode: true);
        }

        private void RunProcess(string fileName, string[] args, string description, bool ignoreExitCode)
        {
            if (options.DryRun)
            {
                Log("[dry-run] " + description + ": " + fileName + " " + string.Join(" ", args.Select(QuoteArg)));
                return;
            }

            Log(description);
            using var process = Process.Start(new ProcessStartInfo
            {
                FileName = fileName,
                ArgumentList = { },
                UseShellExecute = false,
                CreateNoWindow = true
            }.WithArgs(args));
            process!.WaitForExit();
            if (!ignoreExitCode && process.ExitCode != 0)
            {
                throw new InvalidOperationException(fileName + " failed with exit code " + process.ExitCode);
            }
        }

        private void RunRegistryAction(string description, Action action)
        {
            if (options.SimulateRegistryAccessDenied)
            {
                Log("[simulated-denied] " + description);
                throw new UnauthorizedAccessException("Simulated registry access denied.");
            }

            if (options.DryRun)
            {
                Log("[dry-run] " + description);
                return;
            }

            Log(description);
            action();
        }

        private void RunFileAction(string description, Action action)
        {
            if (options.DryRun)
            {
                Log("[dry-run] " + description);
                return;
            }

            Log(description);
            action();
        }

        private void Log(string message) => log.WriteLine(message);
    }

    private static void ValidateTargetPath(string targetDir)
    {
        var full = Path.GetFullPath(targetDir).TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar);
        var forbidden = new[]
        {
            Environment.GetFolderPath(Environment.SpecialFolder.Windows),
            Environment.GetFolderPath(Environment.SpecialFolder.System),
            Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles),
            Environment.GetFolderPath(Environment.SpecialFolder.ProgramFilesX86),
            Environment.GetFolderPath(Environment.SpecialFolder.UserProfile)
        }
        .Where(p => !string.IsNullOrWhiteSpace(p))
        .Select(p => Path.GetFullPath(p).TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar));

        foreach (var path in forbidden)
        {
            if (full.Equals(path, StringComparison.OrdinalIgnoreCase))
            {
                throw new InvalidOperationException("Refusing protected path: " + full);
            }
        }
    }

    private sealed record RegistryTarget(RegistryHive Hive, RegistryView View, string SubKey)
    {
        public override string ToString()
        {
            var hive = Hive == RegistryHive.LocalMachine ? "HKLM" : "HKCU";
            return hive + "\\" + SubKey;
        }
    }

    private sealed record Options
    {
        public bool DryRun { get; private init; }
        public bool Pause { get; private init; }
        public bool Elevate { get; private init; } = true;
        public bool ShowHelp { get; private init; }
        public bool SimulateRegistryAccessDenied { get; private init; }
        public string TargetDir { get; private init; } = DefaultTargetDir;
        public string ChromeUserDataDir { get; private init; } = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), @"Google\Chrome\User Data");
        public string LogPath { get; private init; } = BuildDefaultLogPath();

        public static Options Parse(string[] args)
        {
            var options = new Options();
            for (var i = 0; i < args.Length; i++)
            {
                var arg = args[i];
                switch (arg.ToLowerInvariant())
                {
                    case "--help":
                    case "/?":
                    case "-h":
                        options = options with { ShowHelp = true };
                        break;
                    case "--dry-run":
                    case "/dry-run":
                        options = options with { DryRun = true };
                        break;
                    case "--pause":
                        options = options with { Pause = true };
                        break;
                    case "--no-pause":
                        options = options with { Pause = false };
                        break;
                    case "--elevate":
                        options = options with { Elevate = true };
                        break;
                    case "--no-elevate":
                        options = options with { Elevate = false };
                        break;
                    case "--simulate-registry-access-denied":
                        options = options with { SimulateRegistryAccessDenied = true };
                        break;
                    case "--target-dir":
                        options = options with { TargetDir = ReadValue(args, ref i, arg) };
                        break;
                    case "--chrome-user-data":
                        options = options with { ChromeUserDataDir = ReadValue(args, ref i, arg) };
                        break;
                    case "--log":
                        options = options with { LogPath = ReadValue(args, ref i, arg) };
                        break;
                    default:
                        throw new ArgumentException("Unknown option: " + arg);
                }
            }

            return options;
        }

        private static string ReadValue(string[] args, ref int index, string name)
        {
            if (index + 1 >= args.Length || string.IsNullOrWhiteSpace(args[index + 1]))
            {
                throw new ArgumentException(name + " requires a value.");
            }

            index++;
            return args[index];
        }

        public static void PrintHelp()
        {
            Console.WriteLine("Usage: ONE.CleanupAdmin.exe [--dry-run] [--no-pause] [--no-elevate]");
            Console.WriteLine("       [--target-dir PATH] [--chrome-user-data PATH] [--log PATH]");
            Console.WriteLine("       [--simulate-registry-access-denied]");
        }

        private static string BuildDefaultLogPath()
        {
            var root = Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData);
            if (string.IsNullOrWhiteSpace(root))
            {
                root = Path.GetTempPath();
            }

            return Path.Combine(root, "ONE", "CleanupAdmin", "cleanup-admin-" + DateTime.Now.ToString("yyyyMMdd-HHmmss") + ".log");
        }
    }

    private sealed class LogWriter : IDisposable
    {
        private readonly StreamWriter? writer;

        private LogWriter(string? filePath, StreamWriter? writer)
        {
            FilePath = filePath;
            this.writer = writer;
        }

        public string? FilePath { get; }

        public static LogWriter Create(string filePath)
        {
            if (string.IsNullOrWhiteSpace(filePath))
            {
                return new LogWriter(null, null);
            }

            try
            {
                var fullPath = Path.GetFullPath(Environment.ExpandEnvironmentVariables(filePath));
                var directory = Path.GetDirectoryName(fullPath);
                if (!string.IsNullOrWhiteSpace(directory))
                {
                    Directory.CreateDirectory(directory);
                }

                var stream = new FileStream(fullPath, FileMode.Append, FileAccess.Write, FileShare.ReadWrite);
                var writer = new StreamWriter(stream, Encoding.UTF8) { AutoFlush = true };
                writer.WriteLine("==== ONE CleanupAdmin " + DateTimeOffset.Now.ToString("u") + " ====");
                return new LogWriter(fullPath, writer);
            }
            catch (Exception ex) when (ex is IOException or UnauthorizedAccessException or ArgumentException or NotSupportedException)
            {
                Console.Error.WriteLine("WARN: log disabled: " + ex.Message);
                return new LogWriter(null, null);
            }
        }

        public void WriteLine(string message = "")
        {
            Console.WriteLine(message);
            if (writer == null)
            {
                return;
            }

            writer.WriteLine(DateTimeOffset.Now.ToString("u") + " " + message);
        }

        public void WriteDetail(string message)
        {
            if (writer == null)
            {
                Console.Error.WriteLine(message);
                return;
            }

            writer.WriteLine(DateTimeOffset.Now.ToString("u") + " " + message);
        }

        public void Dispose()
        {
            writer?.Dispose();
        }
    }
}

internal static class ProcessStartInfoExtensions
{
    public static ProcessStartInfo WithArgs(this ProcessStartInfo startInfo, IEnumerable<string> args)
    {
        foreach (var arg in args)
        {
            startInfo.ArgumentList.Add(arg);
        }

        return startInfo;
    }
}
