package rs

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"io"

	"git.daebt.dev/auth/algo"
)

const (
	KeyRSA         algo.KeyType       = "RSA"
	AlgorithmRS256 algo.AlgorithmType = "RS256"
	AlgorithmRS384 algo.AlgorithmType = "RS384"
	AlgorithmRS512 algo.AlgorithmType = "RS512"
)

type Algo struct {
	k algo.KeyType
	a algo.AlgorithmType
	h crypto.Hash
	p *rsa.PrivateKey
}

func (a *Algo) Sign(payload []byte) ([]byte, error) {
	val, err := a.hasher(payload)
	if err != nil {
		return nil, err
	}

	return rsa.SignPKCS1v15(rand.Reader, a.p, a.h, val)
}

func (a *Algo) Verify(payload, signature []byte) error {
	val, err := a.hasher(payload)
	if err != nil {
		return err
	}

	if rsa.VerifyPKCS1v15(&a.p.PublicKey, a.h, val, signature) != nil {
		return algo.ErrInvalidSignature
	}

	return nil
}

func (a *Algo) Key() algo.KeyType {
	return a.k
}

func (a *Algo) Algo() algo.AlgorithmType {
	return a.a
}

func (a *Algo) hasher(payload []byte) ([]byte, error) {
	val := a.h.New()

	if _, err := val.Write(payload); err != nil {
		return nil, err
	}

	return val.Sum(nil), nil
}

func (a *Algo) WriteKeyPEM(w io.Writer) error {
	return pem.Encode(w, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(a.p),
	})
}

func (a *Algo) WritePubPEM(w io.Writer) error {
	return pem.Encode(w, &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: x509.MarshalPKCS1PublicKey(&a.p.PublicKey),
	})
}

func newRS(a *Algo, o ...Option) (*Algo, error) {
	for _, f := range o {
		if err := f(a); err != nil {
			return nil, err
		}
	}

	if a.p == nil {
		return nil, errors.New("key is not initialized")
	}

	return a, nil
}

func NewRS256(o ...Option) (*Algo, error) {
	return newRS(&Algo{
		k: KeyRSA,
		a: AlgorithmRS256,
		h: crypto.SHA256,
	}, o...)
}

func NewRS384(o ...Option) (*Algo, error) {
	return newRS(&Algo{
		k: KeyRSA,
		a: AlgorithmRS384,
		h: crypto.SHA384,
	}, o...)
}

func NewRS512(o ...Option) (*Algo, error) {
	return newRS(&Algo{
		k: KeyRSA,
		a: AlgorithmRS512,
		h: crypto.SHA512,
	}, o...)
}