package cryptoUtils

import (
	"PD1/internal/client"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"log"
	"os"

	"vendor/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 LoadKeyStore(keyStorePath string) KeyStore {

	var privKey rsa.PrivateKey

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

	password := client.AskUserPassword()
	privKeyInterface, cert, caCerts, err := pkcs12.DecodeChain(certFile, password)
	privKey = privKeyInterface.(rsa.PrivateKey)
	if err != nil {
		log.Panicln("PKCS12 key store couldn't be decoded")
	}
	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, err := tls.X509KeyPair(k.cert.Raw, pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(&k.privKey)}))
	if err != nil {
		log.Panicln("Could not load certificate and privkey to TLS")
	}

	//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},
		ClientCAs:    caCertPool,
	}
	return config
}

func (k KeyStore) GetTLSConfigServer() *tls.Config {
	config := k.GetTLSConfig()

	config.ClientAuth = tls.RequireAndVerifyClientCert

	return config
}

func (k KeyStore) GetTLSConfigClient() *tls.Config {
	config := k.GetTLSConfig()

	config.ServerName = "SERVER"

	return config
}

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 {

   return 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
}