using System;
using System.Collections.Generic;
using System.IO;
using System.Security.Cryptography;
using System.Text;
using Tools;
using Tools.Crypto;
using System.Linq;
using System.Threading.Tasks;
using System.Threading;

namespace encaes
{
    enum Mode
    {
        Failed=0,
        Auto,
        Encrypt,
        Decrypt,
        Genkey,
        RegenKey,
    }
    struct Options
    {
        public bool SubstitutePassword;
        public string Password;
        public string KeyFile;
        public string NewKeyFile;
        public string InputFile;
        public string OutputFile;
        public bool InPlace;

    }
    class Program
    {
        static void usage()
        {
            Console.WriteLine("encaes [-p|-P <password>] [<key file>] [-e|-d] <input file> [<output file>]");
            Console.WriteLine("encaes [-p|-P <password>] --keygen <key file>");
            Console.WriteLine("encaes [-p|-P <password>] --regenkey <key file> <new key file>");

            Console.WriteLine("\n-p\tSpecify password for key.");
            Console.WriteLine("-P\tSpecify password instead of key.");
            Console.WriteLine("-e\tForce encrypt.");
            Console.WriteLine("-e\tForce decrypt.");
            Console.WriteLine("--keygen\tGenerate key");
            Console.WriteLine("--regenkey\tRegenerate key");
        }
        static (Mode Mode, Options Options) ParseArgs(Taker<string> args)
        {
            Mode mode = Mode.Failed;
            Options opt = default;
            try
            {

                while (args.TryTake(out string arg0))
                {
                    if (arg0 == "-p" && opt.Password == null)
                    {
                        opt.Password = args.Take();
                        continue;
                    }
                    else if (arg0 == "-P" && opt.Password == null)
                    {
                        opt.Password = args.Take();
                        opt.SubstitutePassword = true;
                        continue;
                    }
                    else if (arg0 == "--keygen" && args.TryTake(out opt.KeyFile))
                    {
                        mode = Mode.Genkey;
                        return (mode, opt);
                    }
                    else if(arg0 == "--regenkey" && args.TryTake(out opt.KeyFile))
                    {
                        mode = Mode.RegenKey;
                        opt.NewKeyFile = args.Take();
                        return (mode, opt);
                    }
                    else if (!opt.SubstitutePassword && opt.KeyFile==null)
                        opt.KeyFile = arg0;
                    else
                    {
                        if (mode == Mode.Failed)
                        {
                            string encmode = arg0;
                            switch (encmode)
                            {
                                case "-e":
                                    mode = Mode.Encrypt;
                                    break;
                                case "-d":
                                    mode = Mode.Decrypt;
                                    break;
                                default:
                                    mode = Mode.Auto;
                                    opt.InputFile = encmode;
                                    opt.InPlace = !args.TryTake(out opt.OutputFile);
                                    return (mode, opt);
                            }
                            opt.InputFile = args.Take();
                            opt.InPlace = !args.TryTake(out opt.OutputFile);
                            return (mode, opt);
                        }
                    }
                }
            }
            catch(IndexOutOfRangeException)
            {
                return (Mode.Failed, default);
            }
            return (mode, opt);
        }
        static async Task Main(string[] args)
        {
            if (args.Length < 2)
            {
                usage();
                return;
            }
            var (mode, opt) = ParseArgs(new Taker<string>(args));

            AESKey getKey()
            {
                try
                {
                    if (opt.SubstitutePassword)
                    {
                        using (var derive = new Rfc2898DeriveBytes(opt.Password, global::encaes.Properties.Resources.globalSalt))
                        {
                            unsafe
                            {
                                return derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                            }
                        }
                    }
                    else
                    {
                        using (var stream = new FileStream(opt.KeyFile, FileMode.Open, FileAccess.Read))
                        {
                            return ReadAesKey(stream, () =>
                            {
                                if (opt.Password == null) throw new InvalidOperationException("Key is password protected");
                                else
                                    return opt.Password;
                            });
                        }
                    }
                }catch(Exception ex)
                {
                    throw new InvalidOperationException("Key loading failed: " + ex.Message);
                }
            }

            switch (mode)
            {
                case Mode.Genkey:
                    try
                    {
                        using (var kf = new FileStream(opt.KeyFile, FileMode.Create))
                        {
                            GenerateKey(kf, opt);
                        }
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine("Error: Keygen failed (" + ex.GetType().Name + "): " + ex.Message);
                    }
                    break;
                case Mode.Encrypt:
                    try
                    {
                        if (opt.InPlace) throw new NotImplementedException("In-place encryption not yet supported.");
                        using (var aes = Aes.Create())
                        {
                            using var cancel = new CancellationTokenSource();

                            void kill(object sender, ConsoleCancelEventArgs ev)
                            {
                                if (!cancel.IsCancellationRequested)
                                {
                                    cancel.Cancel();
                                    ev.Cancel = true;
                                }
                                else Console.WriteLine("Force exit");
                            }

                            Console.CancelKeyPress += kill;
                            try
                            {
                                var key = getKey();
                                try
                                {
                                    key.ToAes(aes);
                                    using (var enc = aes.CreateEncryptor())
                                    {
                                        await using (var readStream = new FileStream(opt.InputFile, FileMode.Open, FileAccess.Read))
                                        {
                                            await using (var outputStream = new FileStream(opt.OutputFile, FileMode.Create))
                                            {
                                                var fh = FileHeader.Create();
                                                await outputStream.WriteValueUnmanagedAsync(fh, cancel.Token);
                                                using (var cs = new CryptoStream(outputStream, enc, CryptoStreamMode.Write))
                                                {
                                                    await readStream.CopyToAsync(cs, cancel.Token);
                                                }
                                            }
                                        }
                                    }
                                }
                                finally { key.ZeroMemory(); }
                            }
                            finally
                            {
                                Console.CancelKeyPress -= kill;
                            }
                        }
                    }catch(Exception ex)
                    {
                        Console.WriteLine("Error: Encryption failed ("+ex.GetType().Name+"): " + ex.Message);
                    }
                    break;
                case Mode.Decrypt:
                    try
                    {
                        if (opt.InPlace) throw new NotImplementedException("In-place encryption not yet supported.");
                        using (var aes = Aes.Create())
                        {
                            using var cancel = new CancellationTokenSource();

                            void kill(object sender, ConsoleCancelEventArgs ev)
                            {
                                if (!cancel.IsCancellationRequested)
                                {
                                    cancel.Cancel();
                                    ev.Cancel = true;
                                }
                                else Console.WriteLine("Force exit");
                            }

                            Console.CancelKeyPress += kill;
                            try
                            {
                                var k = getKey();
                                try
                                {
                                    k.ToAes(aes);

                                    await using (var inputFile = new FileStream(opt.InputFile, FileMode.Open, FileAccess.Read))
                                    {
                                        var fh = await inputFile.ReadValueUnmanagedAsync<FileHeader>(cancel.Token);
                                        if (!fh.Valid) throw new InvalidOperationException("Not an encrypted file");

                                        using (var dec = aes.CreateDecryptor())
                                        {
                                            await using (var outputFile = new FileStream(opt.OutputFile, FileMode.Create))
                                            {
                                                using (var cs = new CryptoStream(inputFile, dec, CryptoStreamMode.Read))
                                                {
                                                    await cs.CopyToAsync(outputFile, cancel.Token);
                                                }
                                            }
                                        }
                                    }
                                }
                                finally { k.ZeroMemory(); }
                            }
                            finally
                            {
                                Console.CancelKeyPress -= kill;
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine("Error: Decryption failed (" + ex.GetType().Name + "): " + ex.Message);
                    }
                    break;
                case Mode.Auto:
                    try
                    {
                        if (opt.InPlace) throw new NotImplementedException("In-place encryption not yet supported.");
                        using (var aes = Aes.Create())
                        {
                            using var cancel = new CancellationTokenSource();

                            void kill(object sender, ConsoleCancelEventArgs ev)
                            {
                                if (!cancel.IsCancellationRequested)
                                {
                                    cancel.Cancel();
                                    ev.Cancel = true;
                                }
                                else Console.WriteLine("Force exit");
                            }

                            Console.CancelKeyPress += kill;
                            try
                            {
                                var k = getKey();
                                //Console.WriteLine("Using key " + BitConverter.ToString(k.ToByteArrayUnmanaged().SHA256Hash()).Replace("-", "").ToLower());
                                try
                                {
                                    k.ToAes(aes);
                                    await using (var inputFile = new FileStream(opt.InputFile, FileMode.Open, FileAccess.Read))
                                    {
                                        FileHeader readh;
                                        if (inputFile.Length < FileHeader.Size || !((readh = await inputFile.ReadValueUnmanagedAsync<FileHeader>(cancel.Token))).Valid)
                                        {
                                            inputFile.Position = 0;
                                            //Encrypting.
                                            using (var enc = aes.CreateEncryptor())
                                            {
                                                await using (var outputFile = new FileStream(opt.OutputFile, FileMode.Create))
                                                {
                                                    var fh = FileHeader.Create();
                                                    await outputFile.WriteValueUnmanagedAsync(fh, cancel.Token);
                                                    using (var cs = new CryptoStream(outputFile, enc, CryptoStreamMode.Write))
                                                    {
                                                        await inputFile.CopyToAsync(cs, cancel.Token);
                                                        cs.Flush();
                                                        if (!cs.HasFlushedFinalBlock)
                                                            cs.FlushFinalBlock();
                                                    }
                                                }
                                            }
                                        }
                                        else
                                        {
                                            //Decrypting.
                                            using (var dec = aes.CreateDecryptor())
                                            {
                                                await using (var outputFile = new FileStream(opt.OutputFile, FileMode.Create))
                                                {
                                                    using (var cs = new CryptoStream(inputFile, dec, CryptoStreamMode.Read))
                                                    {
                                                        await cs.CopyToAsync(outputFile, cancel.Token);
                                                        cs.Flush();
                                                        if (!cs.HasFlushedFinalBlock)
                                                            cs.FlushFinalBlock();
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                                finally { k.ZeroMemory(); }
                            }
                            finally
                            {
                                Console.CancelKeyPress -= kill;
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine("Error: Auto failed (" + ex.GetType().Name + "): " + ex.Message);
                    }
                    break;
                case Mode.RegenKey:
                    try
                    {
                        using var cancel = new CancellationTokenSource();

                        void kill(object sender, ConsoleCancelEventArgs ev)
                        {
                            if (!cancel.IsCancellationRequested)
                            {
                                cancel.Cancel();
                                ev.Cancel = true;
                            }
                            else Console.WriteLine("Force exit");
                        }

                        Console.CancelKeyPress += kill;
                        try
                        {
                            AESKey key;
                            using (var kfs = new FileStream(opt.KeyFile, FileMode.Open, FileAccess.Read))
                            {
                                key = await AesEncryptor.LoadKey(kfs, async () =>
                                {
                                    return await ReadPassword("Key Password: ", cancel.Token);
                                }, cancel.Token);
                            }
                            using (var ofs = new FileStream(opt.NewKeyFile, FileMode.Create))
                            {
                                await AesEncryptor.SaveKey(ofs, key, opt.Password, cancel.Token);
                            }
                        }
                        finally
                        {
                            Console.CancelKeyPress -= kill;
                        }
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine("Error: Regen failed (" + ex.GetType().Name + "): " + ex.Message);
                    }

                        break;
                case Mode.Failed:
                default:
                    usage();
                    return;
            }
        }

        static Task<string> ReadPassword(string prompt = null, CancellationToken token = default)
        {
            TaskCompletionSource<string> comp = new TaskCompletionSource<string>();
            List<char> pwd = new List<char>();

            bool set = false;
            _ = Task.Run(() =>
            {
                if (prompt != null)
                    Console.Write(prompt);
                try
                {
                    using var _c = token.Register(() =>
                    {
                        if (!set)
                        {
                            set = true;
                            comp.SetException(new OperationCanceledException());
                        }
                    });
                    while (true)
                    {
                        ConsoleKeyInfo i = Console.ReadKey(true);
                        if (token.IsCancellationRequested) break;
                        if (i.Key == ConsoleKey.Enter)
                        {
                            Console.WriteLine();
                            break;
                        }
                        else if (i.Key == ConsoleKey.Backspace)
                        {
                            if (pwd.Count > 0)
                            {
                                pwd.RemoveAt(pwd.Count - 1);
                            }
                        }
                        else if (i.KeyChar != '\u0000')
                        {
                            pwd.Add(i.KeyChar);
                        }
                    }

                    if (!set)
                    {
                        set = true;
                        if (token.IsCancellationRequested)
                            comp.SetException(new OperationCanceledException());
                        else
                            comp.SetResult(new string(pwd.ToArray()));
                    }
                }
                catch (Exception ex)
                {
                    Console.WriteLine();
                    pwd.Clear();
                    if (!set)
                    {
                        set = true;
                        comp.SetException(ex);
                    }
                }
            });
            return comp.Task;
        }
        static unsafe AESKey ReadAesKey(Stream from, Func<string> getPasswordFunc)
        {
            var header = from.ReadValueUnmanaged<KeyHeader>();
            if (header.Check != KeyHeader.CheckBit) throw new InvalidDataException("This is not a key");
            if (header.PasswordProtected)
            {
                var passw = getPasswordFunc();
                if (header.PasswordHash != KeyHeader.HashAndSalt(passw, header.Salt.ToByteArrayUnmanaged()))
                    throw new InvalidDataException("Bad password hash");
                using (var derive = new Rfc2898DeriveBytes(passw, header.Salt))
                {
                    var aesk = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                    try
                    {
                        using (var aes = Aes.Create())
                        {
                            aesk.ToAes(aes);
                            using (var dec = aes.CreateDecryptor())
                            {
                                using (var cs = new CryptoStream(from, dec, CryptoStreamMode.Read))
                                {
                                    return cs.ReadValueUnmanaged<AESKey>();
                                }
                            }
                        }
                    }
                    finally
                    {
                        aesk.ZeroMemory();
                    }
                }
            }
            else
            {
                return from.ReadValueUnmanaged<AESKey>();
            }
        }

        static unsafe void GenerateKey(Stream to, Options opt)
        {
            Span<byte> salt = stackalloc byte[PasswordSaltSize];
            using var rng = RandomNumberGenerator.Create();
            if (opt.Password != null)
            {
                rng.GetBytes(salt);
            }

            KeyHeader header;
            AESKey key;
            try
            {
                if (opt.SubstitutePassword)
                {
                    using (var derive = new Rfc2898DeriveBytes(opt.Password, salt.ToArray()))
                    {
                        key = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                    }
                    header = KeyHeader.Create();

                    to.WriteValueUnmanaged(header);
                    to.WriteValueUnmanaged(key);
                }
                else if (opt.Password != null)
                {
                    key = AESKey.NewKey();
                    header = KeyHeader.Create(KeyHeader.HashAndSalt(opt.Password, salt), new Tools.ByValData.Fixed.FixedByteArray16(salt.ToArray()));

                    to.WriteValueUnmanaged(header);
                    using (var derive = new Rfc2898DeriveBytes(opt.Password, salt.ToArray()))
                    {
                        AESKey pwk = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                        using (var aes = Aes.Create())
                        {
                            pwk.ToAes(aes);
                            using (var enc = aes.CreateEncryptor())
                            {
                                using (var cs = new CryptoStream(to, enc, CryptoStreamMode.Write))
                                    cs.WriteValueUnmanaged(key);
                            }
                        }
                    }
                }
                else
                {
                    key = AESKey.NewKey();
                    header = KeyHeader.Create();

                    to.WriteValueUnmanaged(header);
                    to.WriteValueUnmanaged(key);
                }
            }
            finally
            {
                key.ZeroMemory();
            }
        }

        static unsafe byte[] PasswordDecrypt(byte[] buffer, string password, byte[] salt)
        {
            using (var derive = new Rfc2898DeriveBytes(password, salt))
            {
                var aesk = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                try
                {
                    using (var aes = Aes.Create())
                    {
                        aesk.ToAes(aes);
                        using (var dec = aes.CreateDecryptor())
                        {
                            using (var ms = new MemoryStream(buffer))
                            {
                                ms.Position = 0;
                                using (var cs = new CryptoStream(ms, dec, CryptoStreamMode.Read))
                                {
                                    using (var oms = new MemoryStream())
                                    {
                                        cs.CopyTo(oms);
                                        return oms.ToArray();
                                    }
                                }
                            }
                        }
                    }
                }
                finally
                {
                    aesk.ZeroMemory();
                }
            }
        }
        static unsafe byte[] PasswordEncrypt(byte[] buffer, string password, byte[] salt)
        {
            using (var derive = new Rfc2898DeriveBytes(password, salt))
            {
                var aesk = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                try
                {
                    using (var aes = Aes.Create())
                    {
                        aesk.ToAes(aes);
                        using (var enc = aes.CreateEncryptor())
                        {
                            using (var ms = new MemoryStream())
                            {
                                using (var cs = new CryptoStream(ms, enc, CryptoStreamMode.Write))
                                {
                                    cs.Write(buffer, 0, buffer.Length);
                                    cs.Flush();
                                }
                                return ms.ToArray();
                            }
                        }
                    }
                }
                finally
                {
                    aesk.ZeroMemory();
                }
            }
        }

        public const int PasswordSaltSize= 16;
    }

    public unsafe readonly struct FileHeader
    {
        public static readonly int Size = sizeof(FileHeader);

        public const uint CheckBit = 0x2BAD1DEA;

        public readonly uint Check;
        //TODO: Add hash
        public readonly SHA256Hash FileHash;//TODO;

        private FileHeader(uint chk)
        {
            Check = chk;
        }

        public static FileHeader Create()
        {
            return new FileHeader(CheckBit);
        }

        public bool Valid => Check == CheckBit;
    }

    public unsafe readonly struct KeyHeader
    {
        public static readonly int Size = sizeof(KeyHeader);
        public const uint CheckBit = 0xABAD1DEA;

        public readonly uint Check;
        public readonly bool PasswordProtected;
        public readonly Tools.ByValData.Fixed.FixedByteArray16 Salt;
        public readonly SHA256Hash PasswordHash;

        public static SHA256Hash HashAndSalt(string password, Span<byte> salt)
        {
            return Encoding.UTF8.GetBytes(password).Concat(salt.ToArray()).ToArray().SHA256Hash();
        }

        private KeyHeader(uint chk, bool pwd, Tools.ByValData.Fixed.FixedByteArray16 salt, SHA256Hash hash)
        {
            Check = chk;
            PasswordProtected = pwd;
            Salt = salt;
            PasswordHash = hash;
        }

        public static KeyHeader Create(SHA256Hash? PasswordHash=null, Tools.ByValData.Fixed.FixedByteArray16 PasswordSalt=default)
        {
            if (PasswordHash == null)
            {
                return new KeyHeader(CheckBit, false, default, default);
            }
            else
                return new KeyHeader(CheckBit, true, PasswordSalt, PasswordHash.Value);
        }
    }

    sealed class Taker<T>
    {
        private readonly object mutex = new object();
        private readonly IEnumerator<T> iter;
        public Taker(IEnumerable<T> from)
        {
            iter = from.GetEnumerator();
        }

        public bool TryTake(out T value)
        {
            lock (mutex)
            {
                if (iter.MoveNext())
                {
                    value = iter.Current;
                    return true;
                }
            }
            value = default;
            return false;
        }

        public T Take()
        {
            if (!TryTake(out T value))
                throw new IndexOutOfRangeException();
            else
                return value;
        }
    }

}