using Google.Protobuf;
using Org.BouncyCastle.Crypto.Agreement;
using Org.BouncyCastle.Crypto.Digests;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Utilities.Encoders;

using System.Buffers;
using System.Diagnostics;
using System.Text;
using Microsoft.Extensions.Primitives;
using Selectron.Espressif.ProvisionTools.Services.ProvisionService.Security;
using Org.BouncyCastle.Asn1.Ocsp;
using Org.BouncyCastle.Utilities;
using Selectron.Espressif.ProvisionTools.Services.Exceptions;
using Org.BouncyCastle.Crypto;
using Org.BouncyCastle.Crypto.Generators;
using Org.BouncyCastle.Security;
using Org.BouncyCastle.X509;
using Org.BouncyCastle.Asn1.Sec;
using Org.BouncyCastle.Pkcs;

namespace Selectron.Espressif.ProvisionTools.Services.ProvisionService;

public class ProvSessionEndpoint(TransportHttp transportHttp)
{
    public async Task<Status> RequestHandshake0Async()
    {
        var cmd = transportHttp.Security1.FormatSessionHandshake0();
        using var content = new ByteArrayContent(cmd.ToByteArray());


        using var response = await transportHttp.SendPostRequestAsync(Path, content).ConfigureAwait(false);
        response.EnsureSuccessStatusCode();

        var responseContent = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
        var cmdResponse = SessionData.Parser.ParseFrom(responseContent);

        return cmdResponse.ProtoCase switch
        {
            _ => transportHttp.Security1.AcceptHandshake0Response(cmdResponse.Sec1),
        };
    }

    public async Task<Status> RequestHandshake1Async()
    {
        var cmd = transportHttp.Security1.FormatSessionHandshake1();
        using var content = new ByteArrayContent(cmd.ToByteArray());

        using var response = await transportHttp.SendPostRequestAsync(Path, content).ConfigureAwait(false);
        response.EnsureSuccessStatusCode();

        var responseContent = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
        var cmdResponse = SessionData.Parser.ParseFrom(responseContent);

        return cmdResponse.ProtoCase switch
        {
            _ => transportHttp.Security1.AcceptHandshake1Response(cmdResponse.Sec1),
        };
    }

    private static readonly string Path = "prov-session";

    private readonly TransportHttp transportHttp = transportHttp;
}

public class Security1
{
    private X25519PrivateKeyParameters prvKey;
    private X25519PublicKeyParameters pubKey;
    private X25519PublicKeyParameters pubKeyDevice;

    private AesCTR cipher = new AesCTR();

    public SessionData FormatSessionHandshake0()
    {
        var cmd = new SessionData
        {
            SecVer = SecSchemeVersion.SecScheme1,
            Sec1 = new Sec1Payload
            {
                Sc0 = new SessionCmd0(),
                Msg = Sec1MsgType.SessionCommand0
            }
        };
        GenerateKeyPair();
        cmd.Sec1.Sc0.ClientPubkey = ByteString.CopyFrom(pubKey.GetEncoded());
        return cmd;
    }

    public SessionData FormatSessionHandshake1()
    {
        Trace.Assert(clientVerifyData != null);
        Trace.Assert(clientVerifyData.Length > 0);

        var cmd = new SessionData
        {
            SecVer = SecSchemeVersion.SecScheme1,
            Sec1 = new Sec1Payload
            {
                Msg = Sec1MsgType.SessionCommand1,
                Sc1 = new SessionCmd1
                {
                    ClientVerifyData = ByteString.CopyFrom(clientVerifyData),
                }
            }
        };
        return cmd;
    }

    private byte[] clientVerifyData = [];

    public Status AcceptHandshake0Response(Sec1Payload response)
    {
        ArgumentNullException.ThrowIfNull(nameof(response));

        var ans = response.Sr0;
        Trace.TraceInformation("handshake 0 status= {0}", nameof(AcceptHandshake0Response), ans.Status.ToString());

        pubKeyDevice = new X25519PublicKeyParameters(ans.DevicePubkey.ToByteArray());

        //ans.DeviceRandom;
        X25519Agreement exchange = new X25519Agreement();
        exchange.Init(prvKey);

        byte[] sharedKey = new byte[exchange.AgreementSize];
        exchange.CalculateAgreement(pubKeyDevice, sharedKey);


        var pop = ByteString.CopyFromUtf8("abcd1234");
        if (pop.Length > 0)
        {
            var hash = new Sha256Digest();
            hash.BlockUpdate(pop.ToByteArray());
            var digest = new byte[hash.GetDigestSize()];
            hash.DoFinal(digest.AsSpan());
            sharedKey = Utils.HexEncoder.Xor(sharedKey, digest);
        }

        // Device random is the initialization vector
        cipher.Init(sharedKey, ans.DeviceRandom.ToByteArray());
        clientVerifyData = Encode(pubKeyDevice.GetEncoded());

        return ans.Status;
    }

    public Status AcceptHandshake1Response(Sec1Payload response)
    {
        ArgumentNullException.ThrowIfNull(nameof(response));
        var ans = response.Sr1;

        Trace.TraceInformation("handshake 1 status= {0}", ans.Status.ToString());

        byte[] deviceVerifyData = ans.DeviceVerifyData.ToByteArray();
        byte[] decryptedDeviceVerifyData = Decode(deviceVerifyData);
        if (decryptedDeviceVerifyData.SequenceEqual(pubKey.GetEncoded()) == false)
        {
            throw new DeviceSecurityException("Session establishment failed");
        }
        return ans.Status;
    }


    public byte[] Encode(byte[] plainBytes) => cipher.Encode(plainBytes);

    public byte[] Decode(byte[] encryptedBytes) => cipher.Decode(encryptedBytes);


    protected void GenerateKeyPair()
    {
        prvKey = new X25519PrivateKeyParameters(Hex.Decode("a546e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449ac4"), 0);
        pubKey = prvKey.GeneratePublicKey();
        //var publicKeyBob = new X25519PublicKeyParameters(Hex.Decode("e6db6867583030db3594c1a424b15f7c726624ec26b3353b10a903a6d0ab1c4c"), 0);

        //X25519Agreement agreementAlice = new X25519Agreement();
        //agreementAlice.Init(prvKey);
        //byte[] secretAlice = new byte[agreementAlice.AgreementSize];
        //agreementAlice.CalculateAgreement(publicKeyBob, secretAlice, 0);

        //Console.WriteLine(Hex.ToHexString(secretAlice)); // c3da55379de9c6908e94ea4df28d084f32eccf03491c71f754b4075577a28552
    }
}

public class AesCTR
{
    public AesCTR()
    {
        _ = Aes.Create();
        aesCipher = CipherUtilities.GetCipher("AES/CTR/NoPadding");
    }

    public void Init(byte[] messageKey, byte[] ivString)
    {
        var keyParameter = ParameterUtilities.CreateKeyParameter("AES", messageKey);
        cipherParameters = new ParametersWithIV(keyParameter, ivString);
    }

    public byte[] Encode(byte[] plainBytes)
    {
        byte[] outputBytes = new byte[aesCipher.GetOutputSize(plainBytes.Length)];
        aesCipher.Reset();
        aesCipher.Init(true, cipherParameters);
        int length = aesCipher.ProcessBytes(plainBytes, outputBytes, 0);
        aesCipher.DoFinal(outputBytes, length); //Do the final block
        return outputBytes;
    }

    public byte[] Decode(byte[] encryptedBytes)
    {
        byte[] comparisonBytes = new byte[aesCipher.GetOutputSize(encryptedBytes.Length)];
        aesCipher.Reset();
        aesCipher.Init(false, cipherParameters);
        int length = aesCipher.ProcessBytes(encryptedBytes, comparisonBytes, 0);
        aesCipher.DoFinal(comparisonBytes, length); //Do the final block
        return comparisonBytes;
    }

    private readonly IBufferedCipher aesCipher;
    private ICipherParameters? cipherParameters;
}


// command from viewmodel
// public class ... {
    [RelayCommand(CanExecute = nameof(CanUploadAndApplyCredentials))]
    private async Task UploadAndApplyCredentials()
    {
        var ssid = ByteString.CopyFromUtf8(SelectedSsid);
        var passphrase = ByteString.CopyFromUtf8(PassphraseForSelectedSsid);

        var statusProvConfig = await transport.RequestProvConfigAsync(ssid, passphrase);
        if (statusProvConfig != Status.Success)
        {
            if (DialogService != null)
            {
                await DialogService.ShowMessageBoxAsync("Warning", "Provision Configuration", $"status = {statusProvConfig}", icon: "warning", "OK");
            }
            return;
        }

        var statusProvApply = await transport.RequestProvApplyAsync();
        if (statusProvApply != Status.Success)
        {
            if (DialogService != null)
            {
                await DialogService.ShowMessageBoxAsync("Warning", "Provision Apply", $"status = {statusProvApply}", icon: "warning", "OK");
            }
            return;
        }
    }

// }