using Newtonsoft.Json;
using System.Diagnostics.CodeAnalysis;
using System.Security.Cryptography;
using System.Text;

namespace SimpleHttpServer.Login;

internal struct SerialLoginData {
    public string passwordSalt;
    public string extraDataSalt;
    public string pwd;
    public string extraData;

    public LoginData ToPlainData() {
        return new LoginData {
            passwordSalt = Convert.FromBase64String(passwordSalt),
            extraDataSalt = Convert.FromBase64String(extraDataSalt)
        };
    }
}

internal struct LoginData {
    public byte[] passwordSalt;
    public byte[] extraDataSalt;
    public byte[] passwordHash;
    public byte[] encryptedExtraData;

    public SerialLoginData ToSerial() {
        return new SerialLoginData {
            passwordSalt = Convert.ToBase64String(passwordSalt),
            extraDataSalt = Convert.ToBase64String(extraDataSalt),
            pwd = Convert.ToBase64String(passwordHash),
            extraData = Convert.ToBase64String(encryptedExtraData)
        };
    }
}

internal struct LoginDataProviderConfig {

    /// <summary>
    /// Size of the password salt and the extradata salt. So each salt will be of size <see cref="SALT_SIZE"/>.
    /// </summary>
    public int SALT_SIZE = 32;
    public int KEY_LENGTH = 256 / 8;
    public int PBKDF2_ITERATIONS = 600_000;

    public LoginDataProviderConfig() { }
}

public class LoginProvider<TExtraData> {

    private static readonly Func<TExtraData, byte[]> JsonSerialize = t => Encoding.UTF8.GetBytes(JsonConvert.SerializeObject(t));
    private static readonly Func<byte[], TExtraData> JsonDeserialize = b => JsonConvert.DeserializeObject<TExtraData>(Encoding.UTF8.GetString(b))!;

    [ThreadStatic]
    private static SHA256? _sha256PerThread;
    private static SHA256 Sha256PerThread { get => _sha256PerThread ??= SHA256.Create(); }

    private readonly LoginDataProviderConfig config;
    private readonly ReaderWriterLockSlim ldLock = new ReaderWriterLockSlim(LockRecursionPolicy.SupportsRecursion);
    private readonly string ldPath;
    private readonly Dictionary<string, LoginData> loginDatas;

    private Func<TExtraData, byte[]> DataSerializer = JsonSerialize;
    private Func<byte[], TExtraData> DataDeserializer = JsonDeserialize;
    public void SetDataSerializers(Func<TExtraData, byte[]> serializer, Func<byte[], TExtraData> deserializer) {
        DataSerializer = serializer ?? JsonSerialize;
        DataDeserializer = deserializer ?? JsonDeserialize;
    }


    public LoginProvider(string ldPath, string confPath) {
        this.ldPath = ldPath;
        loginDatas = LoadLoginDatas(ldPath);
        config = LoadLoginProviderConfig(confPath);
    }

    private static Dictionary<string, LoginData> LoadLoginDatas(string path) {
        Dictionary<string, SerialLoginData> tempData;
        if (!File.Exists(path)) {
            File.WriteAllText(path, "{}", Encoding.UTF8);
            tempData = new();
        } else {
            tempData = JsonConvert.DeserializeObject<Dictionary<string, SerialLoginData>>(File.ReadAllText(path))!;
            if (tempData == null) {
                throw new InvalidDataException($"could not read login data from file {path}");
            }
        }
        var ld = new Dictionary<string, LoginData>();
        foreach (var pair in tempData) {
            ld.Add(pair.Key, pair.Value.ToPlainData());
        }
        return ld;
    }

    private void SaveLoginData() {
        var serial = new Dictionary<string, SerialLoginData>();
        ldLock.EnterWriteLock();
        try {
            foreach (var pair in loginDatas) {
                serial.Add(pair.Key, pair.Value.ToSerial());
            }
        } finally {
            ldLock.ExitWriteLock();
        }
        File.WriteAllText(ldPath, JsonConvert.SerializeObject(serial));
    }

    private static LoginDataProviderConfig LoadLoginProviderConfig(string path) {
        if (!File.Exists(path)) {
            var conf = new LoginDataProviderConfig();
            File.WriteAllText(path, JsonConvert.SerializeObject(conf));
            return conf;
        }
        return JsonConvert.DeserializeObject<LoginDataProviderConfig>(File.ReadAllText(path));
    }

    public bool AddUser(string username, string password, TExtraData additional) {
        ldLock.EnterWriteLock();
        try {
            if (loginDatas.ContainsKey(username)) {
                return false;
            }
            var passwordSalt = RandomNumberGenerator.GetBytes(config.SALT_SIZE);
            var extraDataSalt = RandomNumberGenerator.GetBytes(config.SALT_SIZE);
            LoginData ld = new LoginData() {
                passwordSalt = passwordSalt,
                extraDataSalt = extraDataSalt,
                passwordHash = ComputeSaltedSha256Hash(password, passwordSalt),
                encryptedExtraData = EncryptExtraData(password, extraDataSalt, additional),
            };
            loginDatas.Add(username, ld);
            SaveLoginData();
        } finally {
            ldLock.ExitWriteLock();
        }
        return true;
    }

    public bool RemoveUser(string username) {
        ldLock.EnterWriteLock();
        try {
            var removed = loginDatas.Remove(username);
            if (removed) {
                SaveLoginData();
            }
            return removed;
        } finally {
            ldLock.ExitWriteLock();
        }
    }

    public bool ModifyUser(string username, string newPassword, TExtraData newExtraData) {
        ldLock.EnterWriteLock();
        try {
            if (!loginDatas.ContainsKey(username)) {
                return false;
            }
            loginDatas.Remove(username, out var data);
            data.passwordHash = ComputeSaltedSha256Hash(newPassword, data.passwordSalt);
            data.encryptedExtraData = EncryptExtraData(newPassword, data.extraDataSalt, newExtraData);
            loginDatas.Add(username, data);
            SaveLoginData();
        } finally {
            ldLock.ExitWriteLock();
        }
        return true;
    }

    public bool TryAuthenticate(string username, string password, [MaybeNullWhen(false)] out TExtraData extraData) {
        LoginData data;
        ldLock.EnterReadLock();
        try {
            if (!loginDatas.TryGetValue(username, out data)) {
                extraData = default;
                return false;
            }
        } finally {
            ldLock.ExitReadLock();
        }
        var hash = ComputeSaltedSha256Hash(password, data.passwordSalt);
        if (!hash.SequenceEqual(data.passwordHash)) {
            extraData = default;
            return false;
        }
        extraData = DecryptExtraData(password, data.extraDataSalt, data.encryptedExtraData);
        return true;
    }

    /// <summary>
    /// Threadsafe as the SHA256 instance (<see cref="Sha256PerThread"/>) is per thread.
    /// </summary>
    /// <param name="data"></param>
    /// <param name="salt"></param>
    /// <returns></returns>
    private static byte[] ComputeSaltedSha256Hash(string data, byte[] salt) {
        var dataBytes = Encoding.UTF8.GetBytes(data);
        var buf = new byte[data.Length + salt.Length];
        Buffer.BlockCopy(dataBytes, 0, buf, 0, dataBytes.Length);
        Buffer.BlockCopy(salt, 0, buf, dataBytes.Length, salt.Length);
        return Sha256PerThread.ComputeHash(buf);
    }

    private byte[] EncryptExtraData(string pwd, byte[] salt, TExtraData extraData) {
        var pbkdf2 = new Rfc2898DeriveBytes(Encoding.UTF8.GetBytes(pwd), salt, config.PBKDF2_ITERATIONS, HashAlgorithmName.SHA256);
        var key = pbkdf2.GetBytes(config.KEY_LENGTH / 8);

        var plainBytes = DataSerializer(extraData);
        using var aes = Aes.Create();
        aes.KeySize = config.KEY_LENGTH;
        aes.Key = key;
        aes.Mode = CipherMode.CBC;
        aes.Padding = PaddingMode.PKCS7;
        ICryptoTransform encryptor = aes.CreateEncryptor(aes.Key, aes.IV);
        byte[] cipherBytes = encryptor.TransformFinalBlock(plainBytes, 0, plainBytes.Length);

        var encryptedBytes = new byte[aes.IV.Length + cipherBytes.Length];
        Array.Copy(aes.IV, 0, encryptedBytes, 0, aes.IV.Length);
        Array.Copy(cipherBytes, 0, encryptedBytes, aes.IV.Length, cipherBytes.Length);

        return encryptedBytes;
    }

    private TExtraData DecryptExtraData(string pwd, byte[] salt, byte[] encryptedData) {
        var pbkdf2 = new Rfc2898DeriveBytes(Encoding.UTF8.GetBytes(pwd), salt, config.PBKDF2_ITERATIONS, HashAlgorithmName.SHA256);
        var key = pbkdf2.GetBytes(config.KEY_LENGTH / 8);

        using var aes = Aes.Create();
        aes.KeySize = config.KEY_LENGTH;
        aes.Key = key;
        aes.Mode = CipherMode.CBC;
        aes.Padding = PaddingMode.PKCS7;
        var iv = new byte[aes.BlockSize / 8];
        var cipherBytes = new byte[encryptedData.Length - iv.Length];

        Array.Copy(encryptedData, 0, iv, 0, iv.Length);
        Array.Copy(encryptedData, iv.Length, cipherBytes, 0, cipherBytes.Length);

        aes.IV = iv;
        ICryptoTransform decryptor = aes.CreateDecryptor(aes.Key, aes.IV);
        byte[] plainBytes = decryptor.TransformFinalBlock(cipherBytes, 0, cipherBytes.Length);

        return DataDeserializer(plainBytes);
    }
}