package certs

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"time"
)

func ToPEM(data []byte, pemType string) ([]byte, error) {
	pemData := new(bytes.Buffer)
	err := pem.Encode(pemData, &pem.Block{
		Type:  pemType,
		Bytes: data,
	})
	if err != nil {
		return nil, err
	}

	return pemData.Bytes(), nil
}

func WriteCert(data []byte, filename string) error {
	// Convert to PEM
	certPEM := new(bytes.Buffer)
	pem.Encode(certPEM, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: data,
	})

	// Write file
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	if _, err := f.Write(certPEM.Bytes()); err != nil {
		return err
	}
	fmt.Printf("Wrote %s.\n", filename)
	return nil
}

func WriteKey(data []byte, filename string) error {
	// Convert to PEM
	KeyPEM := new(bytes.Buffer)
	pem.Encode(KeyPEM, &pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: data,
	})

	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	if _, err := f.Write(KeyPEM.Bytes()); err != nil {
		return err
	}
	fmt.Printf("Wrote %s.\n", filename)
	return nil
}

func GenCACert() (priv []byte, pub []byte, err error) {
	ca := &x509.Certificate{
		SerialNumber: big.NewInt(time.Now().Unix()),
		Subject: pkix.Name{
			Organization: []string{"ezshare"},
			Country:      []string{"NO"},
			Locality:     []string{"Oslo"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour * 24 * 365 * 2),
		IsCA:                  true,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
	}
	caPrivKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	caBytes, err := x509.CreateCertificate(rand.Reader, ca, ca, &caPrivKey.PublicKey, caPrivKey)
	if err != nil {
		return nil, nil, err
	}
	caPrivKeyBytes, err := x509.MarshalECPrivateKey(caPrivKey)
	if err != nil {
		return nil, nil, err
	}

	return caPrivKeyBytes, caBytes, nil
}

func GenAllCerts(path, domain string) error {
	// Create CA certs
	caPriv, caPub, err := GenCACert()
	if err != nil {
		return err
	}
	if err := WriteKey(caPriv, filepath.Join(path, "ca.key")); err != nil {
		return err
	}
	if err := WriteCert(caPub, filepath.Join(path, "ca.pem")); err != nil {
		return err
	}

	// Create server certs
	dnsNames := []string{domain}
	srvKey, srvCrt, err := GenCert("server", caPub, caPriv, dnsNames)
	if err != nil {
		return err
	}
	if err := WriteKey(srvKey, filepath.Join(path, "srv.key")); err != nil {
		return err
	}
	if err := WriteCert(srvCrt, filepath.Join(path, "srv.pem")); err != nil {
		return err
	}

	clientKey, clientCrt, err := GenCert("client", caPub, caPriv, []string{})
	if err != nil {
		return err
	}
	if err := WriteKey(clientKey, filepath.Join(path, "client.key")); err != nil {
		return err
	}
	if err := WriteCert(clientCrt, filepath.Join(path, "client.pem")); err != nil {
		return err
	}

	return nil
}

func GenCert(cn string, caPub, caPrivKey []byte, dnsNames []string) (priv, pub []byte, err error) {
	// Parse ca
	ca, err := x509.ParseCertificate(caPub)
	if err != nil {
		return nil, nil, err
	}

	caPrivKeyParsed, err := x509.ParseECPrivateKey(caPrivKey)
	if err != nil {
		return nil, nil, err
	}

	cert := &x509.Certificate{
		SerialNumber: big.NewInt(time.Now().Unix()),
		Subject: pkix.Name{
			CommonName:   cn,
			Organization: []string{"ezshare"},
			Country:      []string{"No"},
			Locality:     []string{"Oslo"},
		},
		NotBefore:    time.Now(),
		NotAfter:     time.Now().AddDate(10, 0, 0),
		SubjectKeyId: []byte{1, 2, 3, 4, 6},
		DNSNames:     dnsNames,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:     x509.KeyUsageDigitalSignature,
	}

	certPrivKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	certPrivKeyBytes, err := x509.MarshalECPrivateKey(certPrivKey)
	if err != nil {
		return nil, nil, err
	}

	certBytes, err := x509.CreateCertificate(rand.Reader, cert, ca, &certPrivKey.PublicKey, caPrivKeyParsed)
	if err != nil {
		return nil, nil, err
	}
	return certPrivKeyBytes, certBytes, nil
}