package tlsutils

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"golang.org/x/exp/slices"
	"math"
	"math/big"
	"time"
)

func CreateSelfSignedCertificate(name ...string) (*tls.Certificate, error) {
	serial, _ := rand.Int(rand.Reader, big.NewInt(math.MaxInt64))
	x509Template := x509.Certificate{
		SerialNumber:          serial,
		Subject:               pkix.Name{},
		NotBefore:             time.Unix(0, 0).UTC(),
		NotAfter:              time.Date(time.Now().UTC().Year()+10, 0, 0, 0, 0, 0, 0, time.UTC),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if len(name) > 0 {
		x509Template.Subject.CommonName = name[0]
	}

	privateBogusKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	certBytes, err := x509.CreateCertificate(rand.Reader, &x509Template, &x509Template, privateBogusKey.Public(), privateBogusKey)
	if err != nil {
		return nil, err
	}

	keyBytes, err := x509.MarshalPKCS8PrivateKey(privateBogusKey)
	if err != nil {
		return nil, err
	}

	certificate, err := tls.X509KeyPair(pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	}), pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: keyBytes,
	}))

	if err != nil {
		return nil, err
	}

	return &certificate, nil
}

type Configuration struct {
	Config     *tls.Config
	QUICConfig *tls.Config
}

func NewTLSConfiguration(certificatePath, keypairPath, sni string) (*Configuration, error) {
	bogusCertificate, err := CreateSelfSignedCertificate()
	if bogusCertificate == nil {
		return nil, err
	}
	processCertificate(bogusCertificate)

	var serverCertificate tls.Certificate
	if certificatePath != "" && keypairPath != "" {
		serverCertificate, err = tls.LoadX509KeyPair(certificatePath, keypairPath)
		if err != nil {
			return nil, err
		}
		processCertificate(&serverCertificate)
	} else {
		serverCertificate = *bogusCertificate
	}

	tlsConfig := &tls.Config{
		MinVersion: tls.VersionTLS12,
		MaxVersion: 0, //max supported, currently TLS 1.3
		CurvePreferences: []tls.CurveID{
			tls.X25519,
			tls.CurveP256,
			tls.CurveP384,
		},
		CipherSuites: []uint16{
			tls.TLS_CHACHA20_POLY1305_SHA256,
			tls.TLS_AES_256_GCM_SHA384,
			tls.TLS_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
		},
		SessionTicketsDisabled: false,
		NextProtos: []string{
			"http/1.1",
		},
		GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if len(sni) == 0 || sni == info.ServerName {
				return &serverCertificate, nil
			}
			return bogusCertificate, nil
		},
	}

	if serverCertificate.Leaf.PublicKeyAlgorithm == x509.RSA || bogusCertificate.Leaf.PublicKeyAlgorithm == x509.RSA {
		tlsConfig.CipherSuites = append(tlsConfig.CipherSuites, []uint16{
			tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
		}...)
	}

	tlsConfig.GetConfigForClient = func(info *tls.ClientHelloInfo) (*tls.Config, error) {
		configClone := tlsConfig.Clone()

		//Have proper server preference
		for _, suite := range configClone.CipherSuites {
			if slices.Contains(info.CipherSuites, suite) {
				configClone.CipherSuites = []uint16{suite}
				break
			}
		}

		//Have proper server preference
		for _, curve := range configClone.CurvePreferences {
			if slices.Contains(info.SupportedCurves, curve) {
				configClone.CurvePreferences = []tls.CurveID{curve}
				break
			}
		}

		return configClone, nil
	}

	quicTlsConfig := &tls.Config{
		GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if len(sni) == 0 || sni == info.ServerName {
				return &serverCertificate, nil
			}
			return bogusCertificate, nil
		},
	}

	quicTlsConfig.GetConfigForClient = func(info *tls.ClientHelloInfo) (*tls.Config, error) {
		configClone := quicTlsConfig.Clone()

		//Have proper server suite preference
		for _, suite := range configClone.CipherSuites {
			if slices.Contains(info.CipherSuites, suite) {
				configClone.CipherSuites = []uint16{suite}
				break
			}
		}

		//Have proper server curve preference
		for _, curve := range configClone.CurvePreferences {
			if slices.Contains(info.SupportedCurves, curve) {
				configClone.CurvePreferences = []tls.CurveID{curve}
				break
			}
		}

		return configClone, nil
	}

	return &Configuration{
		Config:     tlsConfig,
		QUICConfig: quicTlsConfig,
	}, nil
}

func processCertificate(c *tls.Certificate) {
	if c.Leaf == nil {
		leaf, err := x509.ParseCertificate(c.Certificate[0])
		if err != nil {
			return
		}
		c.Leaf = leaf
	}
}