package auth

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"io/ioutil"
	"os"

	"github.com/golang-jwt/jwt"
)

type KeyStore struct {
	PublicKey  *rsa.PublicKey
	PrivateKey *rsa.PrivateKey
}

func fileExist(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func (ks *KeyStore) LoadPublicKeyFromFile(path string) error {
	pemBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	key, err := jwt.ParseRSAPublicKeyFromPEM(pemBytes)
	if err != nil {
		return err
	}

	ks.PublicKey = key

	return nil
}

func (ks *KeyStore) LoadPublicKeyFromString(str string) error {
	pemBytes, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return err
	}

	key, err := jwt.ParseRSAPublicKeyFromPEM(pemBytes)
	if err != nil {
		return err
	}

	ks.PublicKey = key

	return nil
}

func (ks *KeyStore) LoadPrivateKey(path string) error {
	pemBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	key, err := jwt.ParseRSAPrivateKeyFromPEM(pemBytes)
	if err != nil {
		return err
	}

	ks.PrivateKey = key

	return nil
}

func (ks *KeyStore) GenerateKeys() error {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	ks.PrivateKey = key
	ks.PublicKey = &key.PublicKey

	return nil
}

func (ks *KeyStore) SavePrivateKey(path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(ks.PrivateKey),
	}

	return pem.Encode(file, block)
}

func (ks *KeyStore) SavePublicKey(path string) error {
	bytes, err := x509.MarshalPKIXPublicKey(ks.PublicKey)
	if err != nil {
		return err
	}

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

	block := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: bytes,
	}

	return pem.Encode(file, block)
}

func LoadOrGenerateKeys(privPath, pubPath string) (*KeyStore, error) {
	ks := &KeyStore{}

	if fileExist(privPath) {
		if err := ks.LoadPrivateKey(privPath); err != nil {
			return ks, err
		}
	} else {
		if err := ks.GenerateKeys(); err != nil {
			return ks, err
		}
		if err := ks.SavePrivateKey(privPath); err != nil {
			return ks, err
		}
	}

	if fileExist(pubPath) {
		if err := ks.LoadPublicKeyFromFile(pubPath); err != nil {
			return ks, err
		}
	} else {
		if ks.PublicKey == nil {
			ks.PublicKey = &ks.PrivateKey.PublicKey
		}
		if err := ks.SavePublicKey(pubPath); err != nil {
			return ks, err
		}
	}

	return ks, nil
}