using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Security.Principal;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Forms;
using amecs;
using Ameliorated.ConsoleUtils;
using Microsoft.Win32;

namespace ame_upgrade_preparation_tool
{
    internal class Program
    {
        private static void RunAsUser(Action action)
        {
            var token = NSudo.GetUserToken();
            Task.Run((Action)Delegate.Combine((Action)(() => { NSudo.GetUserPrivilege(token); }),
                action)).Wait();
            Marshal.FreeHGlobal(token);
        }
        public static async Task Main(string[] args)
        {
            ConsoleTUI.Initialize("AME Upgrade Script");

            Console.WriteLine("\r\nChecking drives...");
            
            var drives = DriveInfo.GetDrives();

            List<string> driveLetters = new List<string>();
            foreach (var drive in drives)
            {
                try
                {
                    if (Directory.GetFiles(drive.RootDirectory.FullName).Contains(Path.Combine(drive.RootDirectory.FullName, "setup.exe")) && Directory.GetDirectories(drive.RootDirectory.FullName).Contains(Path.Combine(drive.RootDirectory.FullName, "sources")))
                    {
                        try
                        {
                            string extension = Directory.GetFiles(Path.Combine(drive.RootDirectory.FullName, "sources")).Contains(Path.Combine(drive.RootDirectory.FullName, "sources", "install.wim")) ? ".wim" : Directory.GetFiles(Path.Combine(drive.RootDirectory.FullName, "sources")).Contains(Path.Combine(drive.RootDirectory.FullName, "sources", "install.esd")) ? ".esd" : null;
                            
                            if (string.IsNullOrEmpty(extension))
                                continue;
                            /*
                            var startInfo = new ProcessStartInfo
                            {
                                CreateNoWindow = true,
                                UseShellExecute = false,
                                WindowStyle = ProcessWindowStyle.Normal,
                                RedirectStandardError = true,
                                RedirectStandardOutput = true,
                                FileName = "DISM",
                                Arguments = "/Get-WimInfo /WimFile:\"" + Path.Combine(drive.RootDirectory.FullName, "sources", "install" + extension) + '"'
                            };
                            
                            var exeProcess = new Process
                            {
                                StartInfo = startInfo,
                                EnableRaisingEvents = true
                            };
                            string versionString = null;
                            bool Win11 = false;
                            exeProcess.OutputDataReceived += delegate(object sender, DataReceivedEventArgs outLine)
                            {
                                if (!String.IsNullOrEmpty(outLine.Data))
                                {
                                    var outputString = outLine.Data;
                                    if (outputString.Contains("Version: "))
                                    {
                                        versionString = outputString.Substring(outputString.IndexOf(' ') + 1);
                                    }
                                    else if (outputString.Contains("Windows 11"))
                                    {
                                        Win11 = true;
                                    }
                                }
                            };
                            exeProcess.Start();
                            
                            exeProcess.BeginOutputReadLine();
                            bool exited = exeProcess.WaitForExit(5000);
                            
                            exeProcess.CancelOutputRead();

                            if (versionString == null)
                            {
                                await Task.Delay(500);
                            }
                            */
                            driveLetters.Add(drive.RootDirectory.FullName);
                        } catch (Exception e)
                        {
                            driveLetters.Add(drive.RootDirectory.FullName);
                        }
                    }
                } catch
                {
                    
                }
            }

            if (driveLetters.Count < 1)
            {
                Console.WriteLine("No mounted Windows ISO was detected.\r\n");
                new ChoicePrompt() { Text = "Press any key to Exit...", AnyKey = true}.Start();
                Environment.Exit(0);
            }
            
            if (driveLetters.Count > 1)
            {
                Console.WriteLine("Multiple ISOs detected, please dismount one.\r\n");
                new ChoicePrompt() { Text = "Press any key to Exit...", AnyKey = true}.Start();
                Environment.Exit(0);
            }
            
            var choice = new ChoicePrompt() { AllowEscape = false, Text = "\r\nThis will cause partial de-amelioration. Continue? (Y/N): " }.Start();

            if (choice.Value == 2)
                Environment.Exit(0);
            
            string Username = null;
            string UserDomain = null;
            string UserSID = null;
            try
            {
                NSudo.GetSystemPrivilege();
                
                RunAsUser(() =>
                {
                    Username = WindowsIdentity.GetCurrent().Name.Split('\\').Last();
                    UserDomain = WindowsIdentity.GetCurrent().Name.Split('\\').FirstOrDefault();
                    UserSID = WindowsIdentity.GetCurrent().User.ToString();
                });
            } catch (Exception e)
            {
            }

            bool isSystem = WindowsIdentity.GetCurrent().IsSystem;
            
            try
            {
                string ID = null;
                
                var key = Registry.LocalMachine.OpenSubKey(@"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall");
                foreach (var item in key.GetSubKeyNames())
                {
                    try
                    {
                        if (((string)key.OpenSubKey(item).GetValue("DisplayName")).Equals("Open-Shell"))
                        {
                            ID = item;
                        }
                    } catch (Exception e)
                    {

                    }
                }

                if (ID != null)
                {
                    Console.WriteLine("\r\nUninstalling Open-Shell...");

                    var proc = Process.Start("MsiExec.exe", $"/X{ID} /quiet");
                    proc.WaitForExit();

                    if (UserSID != null)
                    {
                        try
                        {
                            var appData = (string)Registry.Users.OpenSubKey(UserSID + @"\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders").GetValue("AppData");
                            
                            if (Directory.Exists(Path.Combine(appData, "OpenShell")))
                                Directory.Delete(Path.Combine(appData, "OpenShell"), true);
                        } catch (Exception e)
                        {
                            Console.WriteLine("Error: " + e);
                        }
                    }
                    await Task.Delay(5000);

                    if (!Process.GetProcessesByName("explorer").Any())
                    {
                        if (isSystem)
                            NSudo.RunProcessAsUser(NSudo.GetUserToken(), "explorer.exe", "");
                        else
                            Process.Start("explorer.exe");
                    }
                }
            } catch (Exception e)
            {
                Console.WriteLine("Error: " + e.Message);
            }
            
            
            Console.WriteLine("\r\nReverting settings...");
            try
            {
                var polKey = Registry.LocalMachine.OpenSubKey(@"SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\Explorer", true);
                if (polKey.GetValueNames().Contains("SettingsPageVisibility", StringComparer.InvariantCultureIgnoreCase))
                {
                    polKey.DeleteValue("SettingsPageVisibility");
                }

                if (UserSID != null)
                {
                    new Reg.Value() { KeyName = @$"HKU\{UserSID}\SOFTWARE\Policies\Microsoft\Windows\Explorer", ValueName = "DisableNotificationCenter", Operation = Reg.RegistryValueOperation.Delete }.Apply();
                    new Reg.Value() { KeyName = @$"HKU\{UserSID}\SOFTWARE\Classes\CLSID\{{27DD0F8B-3E0E-4ADC-A78A-66047E71ADC5}}\InprocServer32", ValueName = "", Type = Reg.RegistryValueType.REG_SZ, Data = ""}.Apply();
                }
            } catch (Exception e)
            {
                Console.WriteLine("Error: " + e.Message);
            }

            Thread.Sleep(3000);
            Console.WriteLine("\r\nApplying preparation settings...");
            try
            {
                new Reg.Key() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", Operation = RegistryOperation.Add }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", ValueName = "BypassTPMCheck", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", ValueName = "BypassCPUCheck", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", ValueName = "BypassStorageCheck", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", ValueName = "BypassSecureBootCheck", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\LabConfig", ValueName = "BypassRAMCheck", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
                new Reg.Key() { KeyName = @"HKLM\SYSTEM\Setup\MoSetup", Operation = RegistryOperation.Add }.Apply();
                new Reg.Value() { KeyName = @"HKLM\SYSTEM\Setup\MoSetup", ValueName = "AllowUpgradesWithUnsupportedTPMOrCPU", Type = Reg.RegistryValueType.REG_DWORD, Data = 1 }.Apply();
            } catch (Exception e)
            {
                Console.WriteLine("Error: " + e.Message);
            }

            Thread.Sleep(1500);

            var fc = Console.ForegroundColor;
            Console.ForegroundColor = ConsoleColor.Green;
            Console.WriteLine("\r\nCompleted configuration");

            Console.ForegroundColor = fc;

            Console.WriteLine("\r\nStarting Windows Setup...");
            try
            {
                Process.Start(Path.Combine(driveLetters.First(), "setup.exe"), "/Auto Upgrade /DynamicUpdate Disable");
            } catch (Exception e)
            {
                fc = Console.ForegroundColor;
                Console.ForegroundColor = ConsoleColor.Red;
                Console.WriteLine("Error launching Windows Setup: " + e.Message);
                Console.ForegroundColor = fc;
                new ChoicePrompt() { Text = "Press any key to Exit...", AnyKey = true}.Start();
                Environment.Exit(0);
            }
                
            fc = Console.ForegroundColor;
            Console.ForegroundColor = ConsoleColor.Green;
            Console.WriteLine("\r\nSetup launch successful, exiting...");

            Thread.Sleep(6000);
        }
    }
}