using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using System.Threading.Tasks;
using System.Threading;
using System.Threading.Channels;

using Tools;
using Tools.Crypto;
using System.Security.Cryptography;

namespace encaes
{
    public sealed class AsyncMutex : IDisposable
    {
        private class DisposeHook : IDisposable
        {
            public event Action OnDispose;
            public DisposeHook(Action ondispose)
            {
                OnDispose = ondispose;
            }
            public DisposeHook() { }
            public void Dispose()
                => OnDispose?.Invoke();
        }
        private readonly SemaphoreSlim sem;
        public AsyncMutex()
        {
            sem = new SemaphoreSlim(1, 1);
        }

        private AsyncMutex(SemaphoreSlim sem)
        {
            this.sem = sem;
        }

        public async Task<IDisposable> AquireAsync(CancellationToken token = default)
        {
            await sem.WaitAsync(token);
            return new DisposeHook(() => sem.Release());
        }

        public IDisposable Aquire()
        {
            sem.Wait();
            return new DisposeHook(() => sem.Release());
        }

        public void Dispose()
        {
            sem.Dispose();
        }

        public static AsyncMutex Semaphore(int count, int max)
        {
            return new AsyncMutex(new SemaphoreSlim(count, max));
        }
        public static AsyncMutex Semaphore(int count)
        {
            return new AsyncMutex(new SemaphoreSlim(count, count));
        }
    }
    public class AesEncryptor : IDisposable
    {
        public 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;
        }

        protected readonly AsyncMutex readmutex = new AsyncMutex();
        protected readonly CancellationTokenSource cancel = new CancellationTokenSource();
        public Stream File { get; }

        public bool KeepAlive { get; set; } = false;

        public void CancelAllOperations()
        {
            if (!cancel.IsCancellationRequested)
                cancel.Cancel();
        }

        public AesEncryptor(Stream file)
        {
            File = file;
        }

        /// <summary>
        /// Check if <see cref="File"/> is encrypted.
        /// </summary>
        public async Task<bool> StatAsync(CancellationToken token = default)
        {
            using var cancel = CancellationTokenSource.CreateLinkedTokenSource(this.cancel.Token, token);

            using (await readmutex.AquireAsync(cancel.Token))
            {
                if (File.Length < FileHeader.Size) return false;

                File.Position = 0;
                return (await File.ReadValueUnmanagedAsync<FileHeader>(cancel.Token)).Valid;
            }
        }

        public static unsafe AESKey GenerateKeyFromPassword(string password)
        {
            using (var derive = new Rfc2898DeriveBytes(password, global::encaes.Properties.Resources.globalSalt))
            {
                return derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
            }
        }
        public static async Task<AESKey> LoadKey(Stream from, Func<string> getPassword, CancellationToken token = default)
        {
            return await LoadKey(from, async () =>
            {
                await Task.Yield();
                return getPassword();
            }, token);
        }
        public static async Task<AESKey> LoadKey(Stream from, Func<Task<string>> getPassword, CancellationToken token = default)
        {
            if (from.Length < KeyHeader.Size)
            {
                throw new InvalidDataException("Not a key: Too small.");
            }
            var kh = await from.ReadValueUnmanagedAsync<KeyHeader>(token);
            if (kh.Check != KeyHeader.CheckBit)
            {
                throw new InvalidDataException("Not a key: Bad header.");
            }
            if (kh.PasswordProtected)
            {
                var passwd = await getPassword();
                var phash = KeyHeader.HashAndSalt(passwd, kh.Salt.ToByteArrayUnmanaged());
                if (phash != kh.PasswordHash)
                    throw new InvalidDataException("Bad password");
                using (var derive = new Rfc2898DeriveBytes(passwd, kh.Salt.ToByteArrayUnmanaged()))
                {
                    AESKey pkey;
                    unsafe
                    {
                        pkey = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                    }
                    using (var aes = Aes.Create())
                    {
                        pkey.ToAes(aes);
                        using (var dec = aes.CreateDecryptor())
                        {
                            using (var cs = new CryptoStream(from, dec, CryptoStreamMode.Read))
                                return await cs.ReadValueUnmanagedAsync<AESKey>(token);
                        }
                    }
                }
            }
            else
                return await from.ReadValueUnmanagedAsync<AESKey>(token);
        }
        public static Task<AESKey> GenerateKey(Stream to, CancellationToken token = default) => GenerateKey(to, null, token);
        public static async Task<AESKey> GenerateKey(Stream to, string passwordProtect, CancellationToken token = default)
        {
            int aesksize;
            unsafe
            {
                aesksize = sizeof(AESKey);
            }
            if (passwordProtect != null)
            {
                byte[] salt = new byte[Program.PasswordSaltSize];
                byte[] key = new byte[aesksize];

                using (var rng = RandomNumberGenerator.Create())
                {
                    rng.GetBytes(key);
                    rng.GetBytes(salt);
                }

                var passwordHash = KeyHeader.HashAndSalt(passwordProtect, salt);

                KeyHeader kh = KeyHeader.Create(passwordHash, new Tools.ByValData.Fixed.FixedByteArray16(salt));
                await to.WriteValueUnmanagedAsync<KeyHeader>(kh, token);

                using (var derive = new Rfc2898DeriveBytes(passwordProtect, salt))
                {
                    AESKey pkey;
                    unsafe
                    {
                        pkey = derive.GetBytes(sizeof(AESKey)).ToUnmanaged<AESKey>();
                    }
                    using (var aes = Aes.Create())
                    {
                        pkey.ToAes(aes);
                        using (var dec = aes.CreateEncryptor())
                        {
                            using (var cs = new CryptoStream(to, dec, CryptoStreamMode.Write))
                                await cs.WriteAllAsync(key, token);
                        }
                    }
                }

                return key.ToUnmanaged<AESKey>();
            }
            else
            {
                AESKey k;
                await to.WriteValueUnmanagedAsync<KeyHeader>(KeyHeader.Create(), token);
                await to.WriteValueUnmanagedAsync<AESKey>(k = AESKey.NewKey(), token);
                return k;
            }
        }

        public AESKey Key { get; set; } = AESKey.NewKey();

        public static Task SaveKey(Stream to, AESKey key, CancellationToken token = default)
            => SaveKey(to, key, null, token);
        public static async Task SaveKey(Stream to, AESKey key, string passwordProtect, CancellationToken token = default)
        {
            if (passwordProtect == null)
            {
                await to.WriteValueUnmanagedAsync<KeyHeader>(KeyHeader.Create(), token);
                await to.WriteValueUnmanagedAsync(key, token);
            }
            else
            {
                var salt = new byte[Program.PasswordSaltSize];
                using (var rng = RandomNumberGenerator.Create())
                {
                    rng.GetBytes(salt);
                }
                var pwhash = KeyHeader.HashAndSalt(passwordProtect, salt);
                var kh = KeyHeader.Create(pwhash, new Tools.ByValData.Fixed.FixedByteArray16(salt));

                await to.WriteValueUnmanagedAsync(kh, token);
                using (var derive = new Rfc2898DeriveBytes(passwordProtect, salt))
                {
                    int ak;
                    unsafe
                    {
                        ak = sizeof(AESKey);
                    }
                    AESKey pkey = derive.GetBytes(ak).ToUnmanaged<AESKey>();

                    using (var aes = Aes.Create())
                    {
                        pkey.ToAes(aes);
                        using (var enc = aes.CreateEncryptor())
                        {
                            using (var cs = new CryptoStream(to, enc, CryptoStreamMode.Write))
                                await cs.WriteValueUnmanagedAsync(key, token);
                        }
                    }
                }
            }
        }

        public async Task Encrypt(Stream to, CancellationToken token=default)
        {
            using var cancel = CancellationTokenSource.CreateLinkedTokenSource(this.cancel.Token, token);
            using var aes = Aes.Create();
            Key.ToAes(aes);

            using (await readmutex.AquireAsync(cancel.Token))
            {
                File.Position = 0;

                var header = FileHeader.Create();
                await to.WriteValueUnmanagedAsync(header, cancel.Token);

                using (var enc = aes.CreateEncryptor())
                {
                    var cs = new CryptoStream(to, enc, CryptoStreamMode.Write);
                        await File.CopyToAsync(cs, cancel.Token);
                    cs.FlushFinalBlock();
                    await cs.FlushAsync();
                }
            }
        }

        public async Task Decrypt(Stream to, CancellationToken token = default)
        {
            using var cancel = CancellationTokenSource.CreateLinkedTokenSource(this.cancel.Token, token);
            using var aes = Aes.Create();
            Key.ToAes(aes);

            using (await readmutex.AquireAsync(cancel.Token))
            {
                File.Position = 0;

                var header = await File.ReadValueUnmanagedAsync<FileHeader>(cancel.Token);

                if (!header.Valid) throw new InvalidDataException("Not an encrypted stream.");

                using (var dec = aes.CreateDecryptor())
                {
                    using (var cs = new CryptoStream(File, dec, CryptoStreamMode.Read))
                        await cs.CopyToAsync(to, cancel.Token);
                }
            }
        }

        public static async Task<bool> IsEncryptedFile(Stream stream, CancellationToken token=default)
        {
            if (stream.Length < FileHeader.Size) return false;
            var fh = await stream.ReadValueUnmanagedAsync<FileHeader>(token);
            return fh.Valid;
        }

        public static async Task<bool> IsEncryptedFile(string filename, CancellationToken token = default)
        {
            await using (var fs = new FileStream(filename, FileMode.Open, FileAccess.Read))
                return await IsEncryptedFile(fs, token);
        }

        #region IDisposable Support
        private bool disposedValue = false; // To detect redundant calls

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    CancelAllOperations();
                    readmutex.Dispose();
                    cancel.Dispose();

                    if (!KeepAlive)
                        File.Dispose();
                }

                Key = default;

                disposedValue = true;
            }
        }

        ~AesEncryptor()
        {
            Dispose(false);
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }
        #endregion
    }
}