package cryptoUtils

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

	//"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 {

	var privKey *rsa.PrivateKey

	certFile, err := os.ReadFile(keyStorePath)
	if err != nil {
		log.Panicln("Provided certificate couldn't be opened")
	}

	privKeyInterface, cert, caCerts, err := pkcs12.DecodeChain(certFile, password)
	if err != nil {
		log.Panicln("PKCS12 key store couldn't be decoded")
	}
	privKey, ok := privKeyInterface.(*rsa.PrivateKey)
	if !ok {
		log.Panicln("Failed to convert private key to RSA private key")
	}

	if err := privKey.Validate(); err != nil {
		log.Panicln("Private key is not valid")
	}
	return KeyStore{cert: cert, caCertChain: caCerts, privKey: privKey}
}

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
    //FIX: SERVER ACCEPTS CONNECTIONS WITH UNMATCHING OR 
    // NO CERTIFICATE, NEEDS TO BE CHANGED SOMEHOW
	tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
	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 {
	// Digital envolope

	// Create a random symmetric key
	dataKey := make([]byte, 32)
	if _, err := rand.Read(dataKey); err != nil {
		log.Panicln("Could not create dataKey properly: ", err)
	}

	cipher, err := chacha20poly1305.New(dataKey)
	if err != nil {
		log.Panicln("Could not create cipher: ", err)
	}

	nonce := make([]byte, cipher.NonceSize(), cipher.NonceSize()+len(content)+cipher.Overhead())
	if _, err = rand.Read(nonce); err != nil {
		log.Panicln("Could not create data nonce properly: ", 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 {
		log.Panicln("Could not create content signature: ", err)
	}
	content = pair(content, signature)
	ciphertext := cipher.Seal(nonce, nonce, content, nil)

	// crypto/rand.Reader is a good source of entropy for randomizing the
	// encryption function.
	receiverPubKey := receiverCert.PublicKey.(*rsa.PublicKey)
	encryptedDataKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, receiverPubKey, dataKey, nil)
	if err != nil {
		log.Panicln("Could not encrypt dataKey: ", err)
	}

	return pair(encryptedDataKey, ciphertext)
}

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

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

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
}