package dexstore

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"time"

	dexstorage "github.com/dexidp/dex/storage"
	"github.com/go-jose/go-jose/v4"
)

var (
	ErrNoKeys      = errors.New("dexstore: no signing keys")
	errInvalidKeys = errors.New("dexstore: invalid keys")
)

type exportedKeys struct {
	SigningKeyPub    *jose.JSONWebKey          `json:"signing_key_pub"`
	VerificationKeys []exportedVerificationKey `json:"verification_keys"`
}

type exportedVerificationKey struct {
	PublicKey *jose.JSONWebKey `json:"public_key"`
	Expiry    *int64           `json:"expiry"`
}

func ExportPublicKeys(ctx context.Context, dex dexstorage.Storage) (json.RawMessage, error) {
	k, err := dex.GetKeys(ctx)
	if errors.Is(err, dexstorage.ErrNotFound) {
		return nil, ErrNoKeys
	}
	if err != nil {
		return nil, fmt.Errorf("dexstore: read keys: %w", err)
	}
	if k.SigningKeyPub == nil {
		return nil, ErrNoKeys
	}
	out := exportedKeys{
		SigningKeyPub:    k.SigningKeyPub,
		VerificationKeys: make([]exportedVerificationKey, 0, len(k.VerificationKeys)),
	}
	for _, v := range k.VerificationKeys {
		expiry := v.Expiry.UnixMicro()
		out.VerificationKeys = append(out.VerificationKeys, exportedVerificationKey{PublicKey: v.PublicKey, Expiry: &expiry})
	}
	if !publicOnly(out) {
		return nil, errors.New("dexstore: refusing to export private key material")
	}
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(out); err != nil {
		return nil, fmt.Errorf("dexstore: encode keys: %w", err)
	}
	return bytes.TrimSuffix(buf.Bytes(), []byte("\n")), nil
}

func publicOnly(k exportedKeys) bool {
	if k.SigningKeyPub == nil || !k.SigningKeyPub.IsPublic() {
		return false
	}
	for _, v := range k.VerificationKeys {
		if v.PublicKey == nil || !v.PublicKey.IsPublic() {
			return false
		}
	}
	return true
}

func ImportPublicKeys(ctx context.Context, dex dexstorage.Storage, doc json.RawMessage, now time.Time, validFor time.Duration, nextRotation time.Time) error {
	in, err := decodeKeys(doc)
	if err != nil {
		return err
	}
	priv, pub, err := mintKey()
	if err != nil {
		return fmt.Errorf("dexstore: mint signing key: %w", err)
	}
	keys := restoredKeys(in, priv, pub, now, validFor, nextRotation)
	return dex.UpdateKeys(ctx, func(dexstorage.Keys) (dexstorage.Keys, error) {
		return keys, nil
	})
}

func restoredKeys(in dexstorage.Keys, priv, pub *jose.JSONWebKey, now time.Time, validFor time.Duration, nextRotation time.Time) dexstorage.Keys {
	out := dexstorage.Keys{
		SigningKey:       priv,
		SigningKeyPub:    pub,
		VerificationKeys: make([]dexstorage.VerificationKey, 0, len(in.VerificationKeys)+1),
		NextRotation:     nextRotation.UTC(),
	}
	for _, v := range in.VerificationKeys {
		if now.After(v.Expiry) {
			continue
		}
		out.VerificationKeys = append(out.VerificationKeys, v)
	}
	out.VerificationKeys = append(out.VerificationKeys, dexstorage.VerificationKey{PublicKey: in.SigningKeyPub, Expiry: now.Add(validFor).UTC()})
	return out
}

func mintKey() (priv, pub *jose.JSONWebKey, err error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}
	b := make([]byte, 20)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		return nil, nil, err
	}
	kid := hex.EncodeToString(b)
	priv = &jose.JSONWebKey{Key: key, KeyID: kid, Algorithm: string(jose.RS256), Use: "sig"}
	pub = &jose.JSONWebKey{Key: key.Public(), KeyID: kid, Algorithm: string(jose.RS256), Use: "sig"}
	return priv, pub, nil
}

func invalidKeys(reason string) error {
	return fmt.Errorf("%w: %s", errInvalidKeys, reason)
}

func publicKey(k *jose.JSONWebKey, what string) error {
	if k == nil {
		return invalidKeys(what + " is missing")
	}
	if !k.IsPublic() || !k.Valid() {
		return invalidKeys(what + " is not a public key")
	}
	if k.KeyID == "" {
		return invalidKeys(what + " has no kid")
	}
	return nil
}

func decodeKeys(doc json.RawMessage) (dexstorage.Keys, error) {
	var in exportedKeys
	dec := json.NewDecoder(bytes.NewReader(doc))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&in); err != nil {
		return dexstorage.Keys{}, fmt.Errorf("%w: %v", errInvalidKeys, err)
	}
	if dec.More() {
		return dexstorage.Keys{}, invalidKeys("trailing data")
	}
	if err := publicKey(in.SigningKeyPub, "signing_key_pub"); err != nil {
		return dexstorage.Keys{}, err
	}
	out := dexstorage.Keys{SigningKeyPub: in.SigningKeyPub}
	for i, v := range in.VerificationKeys {
		if err := publicKey(v.PublicKey, fmt.Sprintf("verification_keys[%d]", i)); err != nil {
			return dexstorage.Keys{}, err
		}
		if v.Expiry == nil {
			return dexstorage.Keys{}, invalidKeys(fmt.Sprintf("verification_keys[%d] has no expiry", i))
		}
		out.VerificationKeys = append(out.VerificationKeys, dexstorage.VerificationKey{PublicKey: v.PublicKey, Expiry: time.UnixMicro(*v.Expiry).UTC()})
	}
	return out, nil
}