package handshake

import (
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"net"

	"koti.casa/numenor-labs/dsfx/pkg/dcrypto"
	"koti.casa/numenor-labs/dsfx/pkg/encoding/vwb"
)

func debugf(msg string, args ...any) {
	fmt.Printf("(debug): "+msg, args...)
}

// An Actor represents an entity that can participate in the handshake process.
// In this model, there are two actors: Alice and Bob. Alice represents the
// client role and Bob represents the server role. Each actor has a ecdsa private
// key used for identification, and a tcp address that they are reachable at.
type Actor interface {
	// Identity represents the identity of the actor. It is used to sign messages
	// on behalf of the actor during the handshake process. The actor may choose
	// to use a well known key here in order to be recognized by other actors.
	// It is also valid to return an ephemeral key here if the actor wishes to
	// remain anonymous with the remote actor.
	IdentityKey() *ecdsa.PrivateKey

	// Address returns the address of the actor. This is used when listening for
	// incoming connections, and when dialing out to other actors.
	Address() *net.TCPAddr
}

// State represents the state of the handshake process. It should be passed to
// and from the various functions in the handshake process.
type State struct {
	EphemeralKey *ecdh.PrivateKey
}

// InitiateHandshake initiates the handshake process between the given actor
// and the remote actor.
func InitiateHandshake(actor Actor, conn io.ReadWriter, rpub *ecdsa.PublicKey) ([]byte, error) {
	// ------------------------------------------------------------------------
	// Step 1: Ephemeral Key Exchange To Server

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

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

	debugf("client: sending dh key\n")
	// Write the actor's public key to the connection.
	err = vwb.Encode(conn, ourDHKeyRaw)
	if err != nil {
		return nil, err
	}

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

	debugf("client: waiting for server's dh key\n")
	// Read the remote actor's public key from the connection.
	remoteDHKeyRaw, err := vwb.Decode(conn)
	if err != nil {
		return nil, err
	}

	debugf("client: importing server's dh key\n")
	// Import the remote actor's public key.
	remoteDHKey, err := dcrypto.ImportDHPublicKey(remoteDHKeyRaw)
	if err != nil {
		return nil, err
	}

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

	debugf("client: exporting public signing key\n")
	// Export the public key of the actor's signing key.
	ourPublicKeyRaw, err := dcrypto.ExportPublicSigningKey(&actor.IdentityKey().PublicKey)
	if err != nil {
		return nil, err
	}

	debugf("client: exporting remote public signing key\n")
	remotePublicKeyRaw, err := dcrypto.ExportPublicSigningKey(rpub)
	if err != nil {
		return nil, err
	}

	debugf("client: creating authentication message\n")
	// 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.
	authMessage, err := buildMessage(ourDHKey.PublicKey(), remoteDHKey)
	if err != nil {
		return nil, err
	}

	debugf("client: signing authentication message\n")
	// Sign the message with the actor's private key.
	signature, err := dcrypto.Sign(actor.IdentityKey(), authMessage)
	if err != nil {
		return nil, err
	}

	debugf("client: encoding %d bytes of public key\n", len(ourPublicKeyRaw))
	debugf("client: encoding %d bytes of signature\n", len(signature))

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

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

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

	debugf("client: encrypting authentication message\n")
	// Encrypt the message with the derived key.
	boxedMsg, err := dcrypto.Encrypt(derivedKey, plaintext)
	if err != nil {
		return nil, err
	}

	debugf("client: sending authentication message\n")
	// Write the boxed message to the connection.
	err = vwb.Encode(conn, boxedMsg)
	if err != nil {
		return nil, err
	}

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

	debugf("client: waiting for server's authentication message\n")
	// Read the authentication message from the connection.
	authMessageRaw, err := vwb.Decode(conn)
	if err != nil {
		return nil, err
	}

	debugf("client: decrypting authentication message\n")
	// Decrypt the authentication message with the derived key.
	plaintext, err = dcrypto.Decrypt(derivedKey, authMessageRaw)
	if err != nil {
		return nil, err
	}

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

	debugf("client: verifying server's signature\n")
	if !dcrypto.Verify(remotePublicKey, authMessage, plaintext) {
		return nil, errors.New("failed to verify server's signature")
	}

	debugf("client: handshake complete\n")
	return sharedSecret, 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(actor Actor, conn io.ReadWriter) ([]byte, error) {
	// ------------------------------------------------------------------------
	// Step 1: Ephemeral Key Exchange From Client

	debugf("server: starting handshake process\n")
	// Read the remote actor's public key from the connection.
	remoteDHKeyRaw, err := vwb.Decode(conn)
	if err != nil {
		return nil, err
	}

	debugf("server: importing client's dh key\n")
	// Import the remote actor's public key.
	remoteDHKey, err := dcrypto.ImportDHPublicKey(remoteDHKeyRaw)
	if err != nil {
		return nil, err
	}

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

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

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

	debugf("server: sending dh key\n")
	// Write the actor's public key to the connection.
	err = vwb.Encode(conn, ourDHKeyRaw)
	if err != nil {
		return nil, err
	}

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

	debugf("server: waiting for client's authentication message\n")
	// Read the authentication message from the connection.
	authMessageRaw, err := vwb.Decode(conn)
	if err != nil {
		return nil, err
	}

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

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

	debugf("server: decrypting authentication message\n")
	plaintext, err := dcrypto.Decrypt(derivedKey, authMessageRaw)
	if err != nil {
		return nil, err
	}

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

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

	debugf("server: verifying client's signature\n")
	// 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.
	authMessage, err := buildMessage(remoteDHKey, ourDHKey.PublicKey())
	if err != nil {
		return nil, err
	}

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

	debugf("server: creating authentication message\n")
	// 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.
	serverSignature, err := dcrypto.Sign(actor.IdentityKey(), authMessage)
	if err != nil {
		return nil, err
	}

	boxedMsg, err := dcrypto.Encrypt(derivedKey, serverSignature)
	if err != nil {
		return nil, err
	}

	debugf("server: sending authentication message\n")
	// Send the server's signature back to the client.
	err = vwb.Encode(conn, boxedMsg)
	if err != nil {
		return nil, err
	}

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

	debugf("server: handshake complete\n")
	return sharedSecret, 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
}