package rs

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"os"
)

type Option func(*Algo) error

// WithGenerateKey генерирует RSA ключ
func WithGenerateKey(bits int) Option {
	return func(a *Algo) error {
		key, err := rsa.GenerateKey(rand.Reader, bits)
		if err != nil {
			return err
		}

		a.p = key

		return nil
	}
}

// WithKey устанавливает приватный ключ
func WithKey(key *rsa.PrivateKey) Option {
	return func(a *Algo) error {
		a.p = key
		return nil
	}
}

// WithPub устанавливает публичный ключ
func WithPub(pub *rsa.PublicKey) Option {
	return func(a *Algo) error {
		if a.p == nil {
			a.p = &rsa.PrivateKey{}
		}
		a.p.PublicKey = *pub
		return nil
	}
}

// WithKeyPemFile загружает приватный ключ из pem-файла
func WithKeyPEMFile(file string) Option {
	return func(a *Algo) error {
		buf, err := os.ReadFile(file)
		if err != nil {
			return err
		}

		block, _ := pem.Decode(buf)
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return err
		}

		return WithKey(key)(a)
	}
}

func WithPEM(pub, key []byte) Option {
	return func(a *Algo) error {
		if len(key) > 0 {
			block, _ := pem.Decode(key)
			val, err := x509.ParsePKCS1PrivateKey(block.Bytes)
			if err != nil {
				return err
			}

			return WithKey(val)(a)
		}

		block, _ := pem.Decode(pub)
		val, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return err
		}

		return WithPub(val)(a)
	}
}

// WithPubFile загружает публичный ключ из pem-файла
func WithPubFile(file string) Option {
	return func(a *Algo) error {
		buf, err := os.ReadFile(file)
		if err != nil {
			return err
		}

		block, _ := pem.Decode(buf)
		pub, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return err
		}

		return WithPub(pub)(a)
	}
}