package dnet

import (
	"bytes"
	"context"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/sha256"
	"errors"
	"io"
	"log/slog"

	"koti.casa/numenor-labs/dsfx/shared/dcrypto"
	"koti.casa/numenor-labs/dsfx/shared/dlog"
)

type Hand struct {
	State       byte
	Identity    *ecdsa.PrivateKey
	AuthMessage []byte
}

func NewHandshake(identity *ecdsa.PrivateKey) *Hand {
	return &Hand{
		State:    0,
		Identity: identity,
	}
}

// Handshake initiates the handshake process between the given actor
// and the remote actor.
func Handshake(
	ctx context.Context,
	conn io.ReadWriteCloser,
	lPrivKey *ecdsa.PrivateKey,
	rPubKey *ecdsa.PublicKey,
) ([]byte, error) {
	logger := dlog.FromContext(ctx).WithGroup("handshake")

	// ------------------------------------------------------------------------
	// Step 1: Ephemeral Key Exchange To Server

	logger.DebugContext(ctx, "creating dh key")
	// Create a new ECDH private key for the actor.
	ourDHKey, err := dcrypto.GenerateDHKey()
	if err != nil {
		return nil, err
	}

	logger.DebugContext(ctx, "exporting dh key")
	// Export the public key of the actor's ECDH private key.
	ourDHKeyRaw, err := dcrypto.ExportDHPublicKey(ourDHKey.PublicKey())
	if err != nil {
		return nil, err
	}

	// Write the actor's public key to the connection.
	logger.DebugContext(ctx, "sending dh key", slog.Int("key.size", len(ourDHKeyRaw)))
	_, err = NewFrame(ourDHKeyRaw).WriteTo(conn)
	if err != nil {
		return nil, err
	}

	// ------------------------------------------------------------------------
	// Step 2: Ephemeral Key Exchange From Server

	// Read the remote actor's public key from the connection.
	logger.DebugContext(ctx, "waiting for server's dh key")
	remoteDHKeyFrame := NewFrame(nil)
	_, err = remoteDHKeyFrame.ReadFrom(conn)
	if err != nil {
		return nil, err
	}

	// Import the remote actor's public key.
	logger.DebugContext(ctx, "importing server's dh key")
	remoteDHKey, err := dcrypto.ImportDHPublicKey(remoteDHKeyFrame.Contents())
	if err != nil {
		return nil, err
	}

	// ------------------------------------------------------------------------
	// Step 3: Client Authentication

	// Export the public key of the actor's signing key.
	logger.DebugContext(ctx, "exporting public signing key")
	ourPublicKeyRaw, err := dcrypto.ExportPublicSigningKey(&lPrivKey.PublicKey)
	if err != nil {
		return nil, err
	}

	logger.DebugContext(ctx, "exporting remote public signing key")
	remotePublicKeyRaw, err := dcrypto.ExportPublicSigningKey(rPubKey)
	if err != nil {
		return nil, err
	}

	// Construct the message that will be signed by the client.
	// This message is formatted as follows:
	// rlt + sha256(ae + be)
	// This binds both the client and server's long term public keys to the
	// ephemeral keys that were exchanged in the previous step. This creates a
	// verifiable link between both parties and the ephemeral keys used to
	// establish the shared secret.
	logger.DebugContext(ctx, "building authentication message")
	authMessage, err := buildMessage(ourDHKey.PublicKey(), remoteDHKey)
	if err != nil {
		return nil, err
	}

	// Sign the message with the actor's private key.
	logger.DebugContext(ctx, "signing authentication message")
	signature, err := dcrypto.Sign(lPrivKey, authMessage)
	if err != nil {
		return nil, err
	}

	// Compute the shared secret between the actor and the remote actor.
	logger.DebugContext(ctx, "computing shared secret")
	sharedSecret, err := dcrypto.ComputeDHSecret(ourDHKey, remoteDHKey)
	if err != nil {
		return nil, err
	}

	// Derive a key from the shared secret using HKDF.
	logger.DebugContext(ctx, "deriving key from shared secret")
	derivedKey, err := dcrypto.Key(sharedSecret, nil)
	if err != nil {
		return nil, err
	}

	plaintext := make([]byte, 0, len(ourPublicKeyRaw)+len(signature))
	plaintext = append(plaintext, ourPublicKeyRaw...)
	plaintext = append(plaintext, signature...)

	// Encrypt the message with the derived key.
	logger.DebugContext(ctx, "encrypting authentication message")
	boxedMsg, err := dcrypto.Encrypt(derivedKey, plaintext)
	if err != nil {
		return nil, err
	}

	// Write the boxed message to the connection.
	logger.DebugContext(ctx, "sending authentication message", slog.Int("message.size", len(boxedMsg)))
	_, err = NewFrame(boxedMsg).WriteTo(conn)
	if err != nil {
		return nil, err
	}

	// ------------------------------------------------------------------------
	// Step 4: Server Authentication

	// Read the authentication message from the connection.
	logger.DebugContext(ctx, "waiting for server's authentication message")
	authMessageFrame := NewFrame(nil)
	n, err := authMessageFrame.ReadFrom(conn)
	if err != nil {
		return nil, err
	}
	logger.DebugContext(ctx, "received authentication message", slog.Int("message.size", int(n)))

	// Decrypt the authentication message with the derived key.
	logger.DebugContext(ctx, "decrypting authentication message")
	plaintext, err = dcrypto.Decrypt(derivedKey, authMessageFrame.Contents())
	if err != nil {
		return nil, err
	}

	// The server authentication is just verifying the signature it created of
	// the client authentication message.
	logger.DebugContext(ctx, "importing server's public signing key")
	remotePublicKey, err := dcrypto.ImportPublicSigningKey(remotePublicKeyRaw)
	if err != nil {
		return nil, err
	}

	logger.DebugContext(ctx, "verifying server's signature")
	if !dcrypto.Verify(remotePublicKey, authMessage, plaintext) {
		return nil, errors.New("failed to verify server's signature")
	}

	// Finally, we need to let the server know that the handshake is complete.
	logger.DebugContext(ctx, "sending handshake complete message")
	handshakeCompleteMsg := []byte{0x01}
	_, err = NewFrame(handshakeCompleteMsg).WriteTo(conn)
	if err != nil {
		return nil, err
	}

	logger.DebugContext(ctx, "handshake complete")
	return derivedKey, nil
}

// AcceptHandshake accepts a handshake from the given actor and connection. It
// returns the shared secret between the actor and the remote actor.
func AcceptHandshake(ctx context.Context, conn io.ReadWriteCloser, lPrivKey *ecdsa.PrivateKey) (*ecdsa.PublicKey, []byte, error) {
	logger := dlog.FromContext(ctx).WithGroup("handshake")

	// ------------------------------------------------------------------------
	// Step 1: Ephemeral Key Exchange From Client

	// Read the remote actor's public key from the connection.
	logger.DebugContext(ctx, "waiting for client's dh key")
	remoteDHKeyFrame := NewFrame(nil)
	_, err := remoteDHKeyFrame.ReadFrom(conn)
	if err != nil {
		return nil, nil, err
	}

	// Import the remote actor's public key.
	logger.DebugContext(ctx, "importing client's dh key")
	remoteDHKey, err := dcrypto.ImportDHPublicKey(remoteDHKeyFrame.Contents())
	if err != nil {
		return nil, nil, err
	}

	// ------------------------------------------------------------------------
	// Step 2: Ephemeral Key Exchange To Client

	// Create a new ECDH private key for the actor.
	logger.DebugContext(ctx, "creating dh key")
	ourDHKey, err := dcrypto.GenerateDHKey()
	if err != nil {
		return nil, nil, err
	}

	// Export the public key of the actor's ECDH private key.
	logger.DebugContext(ctx, "exporting dh key")
	ourDHKeyRaw, err := dcrypto.ExportDHPublicKey(ourDHKey.PublicKey())
	if err != nil {
		return nil, nil, err
	}

	// Write the actor's public key to the connection.
	logger.DebugContext(ctx, "sending dh key", slog.Int("key.size", len(ourDHKeyRaw)))
	_, err = NewFrame(ourDHKeyRaw).WriteTo(conn)
	if err != nil {
		return nil, nil, err
	}

	// ------------------------------------------------------------------------
	// Step 3: Server Authentication

	// Read the authentication message from the connection.
	logger.DebugContext(ctx, "waiting for client's authentication message")
	authMessageFrame := NewFrame(nil)
	n, err := authMessageFrame.ReadFrom(conn)
	if err != nil {
		return nil, nil, err
	}
	logger.DebugContext(ctx, "received authentication message", slog.Int("message.size", int(n)))

	// Decrypt the authentication message with the derived key.
	logger.DebugContext(ctx, "computing shared secret")
	sharedSecret, err := dcrypto.ComputeDHSecret(ourDHKey, remoteDHKey)
	if err != nil {
		return nil, nil, err
	}

	// Derive a key from the shared secret using HKDF.
	logger.DebugContext(ctx, "deriving key from shared secret")
	derivedKey, err := dcrypto.Key(sharedSecret, nil)
	if err != nil {
		return nil, nil, err
	}

	logger.DebugContext(ctx, "decrypting authentication message")
	plaintext, err := dcrypto.Decrypt(derivedKey, authMessageFrame.Contents())
	if err != nil {
		return nil, nil, err
	}

	clientPublicKeyRaw := plaintext[:222]
	signature := plaintext[222:]

	// Verify the client's public key and signature.
	logger.DebugContext(ctx, "importing client's public signing key")
	clientPublicKey, err := dcrypto.ImportPublicSigningKey(clientPublicKeyRaw)
	if err != nil {
		return nil, nil, err
	}

	// Construct the message that was signed by the client.
	// This message is formatted as follows:
	// rlt + sha256(ae + be)
	// This binds both the client and server's long term public keys to the
	// ephemeral keys that were exchanged in the previous step. This creates a
	// verifiable link between both parties and the ephemeral keys used to
	// establish the shared secret.
	logger.DebugContext(ctx, "building authentication message")
	authMessage, err := buildMessage(remoteDHKey, ourDHKey.PublicKey())
	if err != nil {
		return nil, nil, err
	}

	logger.DebugContext(ctx, "verifying client's signature")
	if !dcrypto.Verify(clientPublicKey, authMessage, signature) {
		return nil, nil, errors.New("failed to verify client's signature")
	}

	// Now we need to sign the authentication message with the server's private
	// key. This will be sent back to the client in the next step to authenticate
	// the server to the client.
	logger.DebugContext(ctx, "signing authentication message")
	serverSignature, err := dcrypto.Sign(lPrivKey, authMessage)
	if err != nil {
		return nil, nil, err
	}

	logger.DebugContext(ctx, "encrypting server's signature")
	boxedMsg, err := dcrypto.Encrypt(derivedKey, serverSignature)
	if err != nil {
		return nil, nil, err
	}

	// Send the server's signature back to the client.
	logger.DebugContext(ctx, "sending authentication message", slog.Int("message.size", len(boxedMsg)))
	_, err = NewFrame(boxedMsg).WriteTo(conn)
	if err != nil {
		return nil, nil, err
	}

	logger.DebugContext(ctx, "waiting for handshake complete message")
	// Read the handshake complete message from the client.
	handshakeCompleteFrame := NewFrame(nil)
	_, err = handshakeCompleteFrame.ReadFrom(conn)
	if err != nil {
		return nil, nil, err
	}
	if !bytes.Equal(handshakeCompleteFrame.Contents(), []byte{0x01}) {
		return nil, nil, errors.New("invalid handshake complete message")
	}

	// ------------------------------------------------------------------------
	// Step 4: Client Authentication

	logger.DebugContext(ctx, "handshake complete")
	return clientPublicKey, derivedKey, nil
}

func buildMessage(clientPubKey *ecdh.PublicKey, serverPubKey *ecdh.PublicKey) ([]byte, error) {
	clientPubKeyRaw, err := dcrypto.ExportDHPublicKey(clientPubKey)
	if err != nil {
		return nil, err
	}
	serverPubKeyRaw, err := dcrypto.ExportDHPublicKey(serverPubKey)
	if err != nil {
		return nil, err
	}
	// Construct the message that will be signed by the client.
	// This message is formatted as follows:
	// rlt + sha256(ae + be)
	// This binds both the client and server's long term public keys to the
	// ephemeral keys that were exchanged in the previous step. This creates a
	// verifiable link between both parties and the ephemeral keys used to
	// establish the shared secret.
	message := make([]byte, 0, len(clientPubKeyRaw)+len(serverPubKeyRaw))
	message = append(message, clientPubKeyRaw...)
	message = append(message, serverPubKeyRaw...)

	messageChecksum := sha256.Sum256(message)
	authMessage := make([]byte, 0, len(serverPubKeyRaw)+sha256.Size)
	authMessage = append(authMessage, serverPubKeyRaw...)
	authMessage = append(authMessage, messageChecksum[:]...)

	return authMessage, nil
}