// KEKS/CM -- KEKS-encoded cryptographic messages
// Copyright (C) 2024-2025 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as
// published by the Free Software Foundation, version 3 of the License.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this program.  If not, see <http://www.gnu.org/licenses/>.

package sign

import (
	"bytes"
	"errors"
	"fmt"
	"hash"
	"time"

	"github.com/google/uuid"

	"go.cypherpunks.su/keks"
	"go.cypherpunks.su/keks/cm"
	ed25519blake2b "go.cypherpunks.su/keks/cm/sign/ed25519-blake2b"
	"go.cypherpunks.su/keks/cm/sign/gost"
	"go.cypherpunks.su/keks/cm/sign/slhdsa"
	"go.cypherpunks.su/keks/schema"
)

const (
	KUCA     = "ca"  // CA-capable key usage
	KUSig    = "sig" // Signing-capable key usage
	KUKEM    = "kem" // Key-encapsulation-mechanism key usage
	FPRLen   = 32    // fingerprint's length
	PubMagic = keks.Magic("cm/pub")
)

var (
	ErrSigInvalid = errors.New("signature is invalid")
	ErrBadSigAlgo = errors.New("bad signature algo")
)

// Public key' contents.
type PubData struct {
	KU   map[string]*struct{} `keks:"ku"`
	Sub  map[string]string    `keks:"sub"`
	Crit []map[string]any     `keks:"crit"`
	Pub  []cm.AV              `keks:"pub"`
	Id   []byte               `keks:"id"`
}

// Parse KEKS-encoded data as Signed with the PubData (certificate) contents.
func PubParse(data []byte) (signed *Signed, tail []byte, err error) {
	{
		var magic keks.Magic
		magic, data = keks.StripMagic(data)
		if magic != "" && magic != PubMagic {
			err = errors.New("wrong magic")
			return
		}
	}
	d := keks.NewDecoderFromBytes(data, nil)
	{
		var v any
		v, err = d.Decode()
		if err != nil {
			return
		}
		err = schema.Check("pub", PubSchemas, v)
		if err != nil {
			return
		}
	}
	d = keks.NewDecoderFromBytes(data, nil)
	var sd Signed
	err = d.DecodeStruct(&sd)
	if err != nil {
		return
	}
	tail = d.B
	signed = &sd
	if sd.TBS.T != "pub" {
		err = errors.New("PubParse: wrong load type")
	}
	return
}

// Check if public key has desired ku capability.
func (pub *PubData) Can(ku string) (yes bool) {
	_, yes = pub.KU[ku]
	return
}

// Sign the provided Signed, having PubData payload with the provided
// parent's PubData and prv key. Certification CID will be automatically
// generated UUIDv7. since and till times must not have nanoseconds part.
func (signed *Signed) CertifyWith(
	parent *PubData,
	prv Iface,
	since, till time.Time,
) error {
	cid, err := uuid.NewV7()
	if err != nil {
		return err
	}
	return signed.SignWith(parent, prv, map[string]any{
		"cid": cid,
		"exp": []time.Time{since, till},
	})
}

// Verify signature of signed data. ErrSigInvalid will be returned in
// case of invalid signature.
func (pub *PubData) CheckSignature(algo string, signed, signature []byte) (err error) {
	if !pub.Can(KUSig) || len(pub.Pub) != 1 {
		err = errors.New("pub can not sign")
		return
	}
	key := pub.Pub[0]
	var valid bool
	switch key.A {
	case ed25519blake2b.Ed25519BLAKE2b:
		if algo != ed25519blake2b.Ed25519BLAKE2b {
			return ErrBadSigAlgo
		}
		valid, err = ed25519blake2b.Verify(key.V, signed, signature)
		if !valid {
			err = ErrSigInvalid
		}
	case gost.GOST3410256A, gost.GOST3410512C:
		if algo != key.A {
			return ErrBadSigAlgo
		}
		valid, err = gost.Verify(key.A, key.V, signed, signature)
		if !valid {
			err = ErrSigInvalid
		}
	case slhdsa.SLHDSASHAKE256s:
		if algo != slhdsa.SLHDSASHAKE256s {
			return ErrBadSigAlgo
		}
		valid, err = slhdsa.Verify(key.A, key.V, signed, signature)
		if !valid {
			err = ErrSigInvalid
		}
	default:
		err = errors.New("unsupported signature algorithm")
	}
	return
}

// Verify signature of signed data, by providing prehashed data.
// ErrSigInvalid will be returned in case of invalid signature.
func (pub *PubData) CheckSignaturePrehash(
	algo string,
	prehash, signature []byte,
) (err error) {
	if !pub.Can(KUSig) || len(pub.Pub) != 1 {
		err = errors.New("pub can not sign")
		return
	}
	key := pub.Pub[0]
	var valid bool
	switch key.A {
	case ed25519blake2b.Ed25519BLAKE2b:
		switch algo {
		case ed25519blake2b.Ed25519PhBLAKE2b:
		case ed25519blake2b.Ed25519PhBLAKE2bMerkle:
		default:
			return ErrBadSigAlgo
		}
		valid, err = ed25519blake2b.VerifyPrehash(key.V, prehash, signature)
		if !valid {
			err = ErrSigInvalid
		}
	case gost.GOST3410256A, gost.GOST3410512C:
		switch algo {
		case gost.GOST3410256A:
		case gost.GOST3410256AMerkle:
		case gost.GOST3410512C:
		case gost.GOST3410512CMerkle:
		default:
			return ErrBadSigAlgo
		}
		valid, err = gost.VerifyPrehash(key.A, key.V, prehash, signature)
		if !valid {
			err = ErrSigInvalid
		}
	case slhdsa.SLHDSASHAKE256s:
		switch algo {
		case slhdsa.SLHDSASHAKE256sPh:
		case slhdsa.SLHDSASHAKE256sMerkle:
		default:
			return ErrBadSigAlgo
		}
		valid, err = slhdsa.VerifyPrehash(algo, key.V, prehash, signature)
		if !valid {
			err = ErrSigInvalid
		}
	default:
		err = errors.New("unsupported signature algorithm")
	}
	return
}

// Verify Signed PubData certification signature with provided parent.
// If prehasher is specified, then prehashed signature mode is used.
// Currently only single signature can be verified.
func (signed *Signed) CertificationCheckSignatureFrom(
	parent *PubData,
	prehasher *hash.Hash,
) (err error) {
	if !parent.Can(KUSig) || len(parent.Pub) != 1 {
		err = errors.New("parent can not sign")
		return
	}
	for _, sig := range signed.Sigs {
		if !bytes.Equal(sig.TBS["sid"].([]byte), parent.Id) {
			continue
		}
		if prehasher == nil {
			var tbs bytes.Buffer
			if _, err = keks.Encode(&tbs, signed.Data, nil); err != nil {
				return
			}
			if _, err = keks.Encode(&tbs, signed.TBS, nil); err != nil {
				return
			}
			if _, err = keks.Encode(&tbs, sig.TBS, nil); err != nil {
				return
			}
			return parent.CheckSignature(sig.Sign.A, tbs.Bytes(), sig.Sign.V)
		} else {
			if _, err = keks.Encode(*prehasher, signed.Data, nil); err != nil {
				return
			}
			if _, err = keks.Encode(*prehasher, signed.TBS, nil); err != nil {
				return
			}
			if _, err = keks.Encode(*prehasher, sig.TBS, nil); err != nil {
				return
			}
			return parent.CheckSignaturePrehash(
				sig.Sign.A, (*prehasher).Sum(nil), sig.Sign.V)
		}
	}
	return errors.New("can not find necessary sid")
}

// Get PubData from Signed.
// Returns nil if Signed does not hold it (or it is not yet parsed).
func (signed *Signed) PubData() *PubData {
	if signed.TBS.T != "pub" || signed.Data == nil {
		return nil
	}
	v, ok := (*signed.Data).(map[string]any)
	if !ok {
		return nil
	}
	var pubData PubData
	err := keks.Map2Struct(&pubData, v)
	if err != nil {
		return nil
	}
	return &pubData
}

// Verify signed Signed PubData certification against pubs chain of
// public keys at specified point of time t.
func (signed *Signed) CertificationVerify(pubs []*Signed, t time.Time) (err error) {
	if len(signed.Sigs) == 0 {
		return errors.New("no sigs")
	}

	var sigTBS *SigTBS
	sigTBS, err = signed.Sigs[0].TBSGet()
	if err != nil {
		return
	}
	{
		exp := sigTBS.Exp
		if t.Before(exp[0]) || t.Equal(exp[0]) {
			err = errors.New("pub is not active")
			return
		}
		if t.After(exp[1]) || t.Equal(exp[1]) {
			err = errors.New("pub is expired")
			return
		}
	}
	sid := sigTBS.SID
	if bytes.Equal(sid, signed.PubData().Id) {
		return signed.CertificationCheckSignatureFrom(signed.PubData(), nil)
	}
	type FPR [FPRLen]byte
	idToPub := make(map[FPR]*Signed, len(pubs))
	for _, cer := range pubs {
		pubData := cer.PubData()
		if !pubData.Can(KUSig) || len(pubData.Pub) != 1 {
			err = errors.New("pub can not sign")
			return
		}
		idToPub[FPR(pubData.Id)] = cer
	}
	signer := idToPub[FPR(sid)]
	if signer == nil {
		err = fmt.Errorf("no pub found for sid: %v", sigTBS.SID)
		return
	}
	err = signed.CertificationCheckSignatureFrom(signer.PubData(), nil)
	if err != nil {
		return
	}
	return signer.CertificationVerify(pubs, t)
}
