package cryptoUtils

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"errors"

	"log"
	"os"

	"golang.org/x/crypto/chacha20poly1305"
	"software.sslmate.com/src/go-pkcs12"
)

type KeyStore struct {
	cert        *x509.Certificate
	caCertChain []*x509.Certificate
	privKey     *rsa.PrivateKey
}

func (k KeyStore) GetCert() *x509.Certificate {
	return k.cert
}

func (k KeyStore) GetCACertChain() []*x509.Certificate {
	return k.caCertChain
}

func (k KeyStore) GetPrivKey() *rsa.PrivateKey {
	return k.privKey
}

func ExtractAllOIDValues(cert *x509.Certificate) map[string]string {
	oidValueMap := make(map[string]string)
	for _, name := range cert.Subject.Names {
		oid := name.Type.String()
		value := name.Value.(string)
		oidValueMap[oid] = value
	}
	return oidValueMap
}

func LoadKeyStore(keyStorePath string, password string) (KeyStore, error) {

	var privKey *rsa.PrivateKey

	keystoreBytes, err := os.ReadFile(keyStorePath)
	if err != nil {
		return KeyStore{}, err
	}

	privKeyInterface, cert, caCerts, err := pkcs12.DecodeChain(keystoreBytes, password)
	if err != nil {
		return KeyStore{}, err
	}
	privKey, ok := privKeyInterface.(*rsa.PrivateKey)
	if !ok {
		return KeyStore{}, err
	}

	if err := privKey.Validate(); err != nil {
		return KeyStore{}, err
	}
	return KeyStore{cert: cert, caCertChain: caCerts, privKey: privKey}, err
}

// Check if the cert is signed by the CA and is for the correct user
func (k KeyStore) CheckCert(cert *x509.Certificate, uid string) error {
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	opts := x509.VerifyOptions{
		Roots: caCertPool,
	}
	// Check if the certificate is signed by the specified CA
	_, err := cert.Verify(opts)
	if err != nil {
		log.Println("Certificate not signed by a trusted CA")
		return err
	}

	//Check if the pseudonym field is set to UID
	oidMap := ExtractAllOIDValues(cert)
	if oidMap["2.5.4.65"] != uid {
		log.Println("Certificate does not belong to the message's receiver")
		return err
	}
	return nil
}

func (k *KeyStore) GetTLSConfig() *tls.Config {

	certificate := tls.Certificate{Certificate: [][]byte{k.cert.Raw}, PrivateKey: k.privKey, Leaf: k.cert}

	//Add the CA certificate chain to a CertPool
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	config := &tls.Config{
		Certificates: []tls.Certificate{certificate},
	}
	return config
}

func (k *KeyStore) GetServerTLSConfig() *tls.Config {
	tlsConfig := k.GetTLSConfig()

	//Add the CA certificate chain to a CertPool
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	tlsConfig.ClientCAs = caCertPool
	tlsConfig.ClientAuth = tls.RequireAnyClientCert
	tlsConfig.VerifyPeerCertificate = func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
		// Verify the peer's certificate
		opts := x509.VerifyOptions{
			Roots: caCertPool,
		}
		for _, certBytes := range rawCerts {
			cert, err := x509.ParseCertificate(certBytes)
			if err != nil {
				return err
			}
			// Check if the certificate is signed by the specified CA
			_, err = cert.Verify(opts)
			if err != nil {
				return errors.New("certificate not signed by trusted CA")
			}
		}
		return nil
	}
	return tlsConfig
}

func (k *KeyStore) GetClientTLSConfig() *tls.Config {
	tlsConfig := k.GetTLSConfig()

	//Add the CA certificate chain to a CertPool
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	tlsConfig.RootCAs = caCertPool
	tlsConfig.InsecureSkipVerify = true
	tlsConfig.VerifyPeerCertificate = func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
		// Verify the peer's certificate
		opts := x509.VerifyOptions{
			Roots: caCertPool,
		}
		for _, certBytes := range rawCerts {
			cert, err := x509.ParseCertificate(certBytes)
			if err != nil {
				return err
			}
			oidMap := ExtractAllOIDValues(cert)

			// Check if the certificate is signed by the specified CA
			_, err = cert.Verify(opts)
			if err != nil {
				return errors.New("certificate not signed by trusted CA")
			}

			//Check if the pseudonym field is set to "SERVER"
			if oidMap["2.5.4.65"] != "SERVER" {
				return errors.New("peer isn't the server")
			}
		}
		return nil
	}
	return tlsConfig
}

func (k KeyStore) EncryptMessageContent(receiverCert *x509.Certificate, content []byte) ([]byte, error) {
	// Digital envolope

	// Create a random symmetric key
	dataKey := make([]byte, 32)
	if _, err := rand.Read(dataKey); err != nil {
		return nil, err
	}

	cipher, err := chacha20poly1305.New(dataKey)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, cipher.NonceSize(), cipher.NonceSize()+len(content)+cipher.Overhead())
	if _, err = rand.Read(nonce); err != nil {
		return nil, err
	}

	// sign the message and append the signature
	hashedContent := sha256.Sum256(content)
	signature, err := rsa.SignPKCS1v15(nil, k.privKey, crypto.SHA256, hashedContent[:])
	if err != nil {
		return nil, err
	}
	content = pair(signature, content)
	ciphertext := cipher.Seal(nonce, nonce, content, nil)

	receiverPubKey := receiverCert.PublicKey.(*rsa.PublicKey)
	encryptedDataKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, receiverPubKey, dataKey, nil)
	if err != nil {
		return nil, err
	}
	return pair(encryptedDataKey, ciphertext), nil
}

func (k KeyStore) DecryptMessageContent(senderCert *x509.Certificate, cipherContent []byte) ([]byte, error) {
	encryptedDataKey, encryptedMsg := unPair(cipherContent)
	dataKey, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, k.GetPrivKey(), encryptedDataKey, nil)
	if err != nil {
		return nil, err
	}
	// decrypt ciphertext
	cipher, err := chacha20poly1305.New(dataKey)
	if err != nil {
		return nil, err
	}

	nonce, ciphertext := encryptedMsg[:cipher.NonceSize()], encryptedMsg[cipher.NonceSize():]
	contentAndSig, err := cipher.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, err
	}
	// check signature with sender public key
	signature, content := unPair(contentAndSig)
	hashedContent := sha256.Sum256(content)
	senderKey := senderCert.PublicKey.(*rsa.PublicKey)
	if err := rsa.VerifyPKCS1v15(senderKey, crypto.SHA256, hashedContent[:], signature); err != nil {
		return nil, err
	}
	return content, nil
}

func pair(l []byte, r []byte) []byte {
	length := len(l)
	lenBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(lenBytes, uint16(length))

	lWithLen := append(lenBytes, l...)
	return append(lWithLen, r...)
}

func unPair(pair []byte) ([]byte, []byte) {
	lenBytes := pair[:2]
	pair = pair[2:]
	length := binary.BigEndian.Uint16(lenBytes)
	l := pair[:length]
	r := pair[length:]
	return l, r
}