package localserver

import (
	"context"
	"crypto"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

const PEM_PRIVATE_KEY = "RSA PRIVATE KEY"
const PEM_PUBLIC_KEY = "RSA PUBLIC KEY"
const OPENSSH_PRIVATE_KEY = "OPENSSH PRIVATE KEY"

type EncryptionManager struct {
	PrivKey  *rsa.PrivateKey
	PeerKeys map[string]*rsa.PublicKey
}

func NewEncryptionManager() (eManager *EncryptionManager) {
	eManager = new(EncryptionManager)
	return
}

func (em *EncryptionManager) GenerateKeyPair(ctx context.Context, name string, pwd string) (privKey *rsa.PrivateKey, err error) {
	privKey, err = rsa.GenerateKey(rand.Reader, 1<<12)
	if err != nil {
		return
	}
	block := &pem.Block{
		Type:  PEM_PRIVATE_KEY,
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	}
	if pwd != "" {
		block, err = x509.EncryptPEMBlock(rand.Reader, block.Type, block.Bytes, []byte(pwd), x509.PEMCipherAES256)
		if err != nil {
			return
		}
	}
	f, err := os.Create(filepath.Join("config", name))
	if err != nil {
		return
	}
	_, err = f.Write(pem.EncodeToMemory(block))
	if err != nil {
		return
	}
	pubBlock := &pem.Block{
		Type:  PEM_PUBLIC_KEY,
		Bytes: x509.MarshalPKCS1PublicKey(&privKey.PublicKey),
	}
	if pwd != "" {
		pubBlock, err = x509.EncryptPEMBlock(rand.Reader, pubBlock.Type, pubBlock.Bytes, []byte(pwd), x509.PEMCipherAES256)
		if err != nil {
			return
		}
	}
	pubf, err := os.Create(filepath.Join("config", fmt.Sprintf("%s.pub", name)))
	if err != nil {
		return
	}
	_, err = pubf.Write(pem.EncodeToMemory(pubBlock))
	return
}

func (em *EncryptionManager) LoadPrivKey(privKeyPath string, password string) (err error) {
	f, err := os.Open(privKeyPath)
	if err != nil {
		return
	}
	var buff []byte
	buff, err = io.ReadAll(f)
	if err != nil {
		return
	}
	privPem, _ := pem.Decode(buff)
	var privePemBytes []byte
	if privPem.Type != PEM_PRIVATE_KEY && privPem.Type != OPENSSH_PRIVATE_KEY {
		logger.Println(privPem.Type)
		err = fmt.Errorf("RSA Private key is of wrong type")
		return
	}
	if password != "" {
		privePemBytes, err = x509.DecryptPEMBlock(privPem, []byte(password))
	} else {
		privePemBytes = privPem.Bytes
	}
	var parsedKey interface{}
	if parsedKey, err = x509.ParsePKCS1PrivateKey(privePemBytes); err != nil {
		logger.Printf("error one %v\n", err)
		if parsedKey, err = x509.ParsePKCS8PrivateKey(privePemBytes); err != nil {
			logger.Printf("error on parsing %v\n", err)
			if parsedKey, err = x509.ParseECPrivateKey(privePemBytes); err != nil {
				return
			}
		}

	}
	var privKey *rsa.PrivateKey
	var ok bool
	privKey, ok = parsedKey.(*rsa.PrivateKey)
	if !ok {
		err = fmt.Errorf("Unable to parse this private key")
		return
	}
	em.PrivKey = privKey
	return
}

func (em *EncryptionManager) ParsePrivKey(privateKey string) (privKey *rsa.PrivateKey, err error) {
	key := []byte(privateKey)
	block, _ := pem.Decode(key)
	b := block.Bytes
	privKey, err = x509.ParsePKCS1PrivateKey(b)
	if err != nil {
		return
	}
	return
}

func (em *EncryptionManager) SignRequest(privKey, id string) (sig string, err error) {
	key, err := em.ParsePrivKey(privKey)
	if err != nil {
		return
	}
	msg := []byte(id)
	hashed := sha256.Sum256(msg)
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, hashed[:])
	if err != nil {
		return
	}
	sig = base64.StdEncoding.EncodeToString(signature)
	return
}

func (em *EncryptionManager) SignRequestHMAC(id string) string {
	h := hmac.New(sha256.New, []byte(NodeToken))
	h.Write([]byte(id))
	sha := h.Sum(nil)
	sig := base64.StdEncoding.EncodeToString(sha)
	return sig
}