// 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 slhdsa

import (
	"crypto"
	"errors"
	"hash"
	"io"

	"go.cypherpunks.su/keks/cm/sign/slhdsa/circl"

	cmhash "go.cypherpunks.su/keks/cm/hash"
	"go.cypherpunks.su/keks/cm/hash/merkle"
	"go.cypherpunks.su/keks/cm/sign/mode"
)

type Signer struct {
	NewHasher func() hash.Hash
	prehasher *hash.Hash
	PK        circl.PublicKey
	SK        circl.PrivateKey
	mode      mode.Mode
}

func (s *Signer) SetMode(m mode.Mode) error {
	switch m {
	case mode.Pure:
		s.mode = m
		return nil
	case mode.Prehash:
		s.mode = m
		p := cmhash.NewSHAKE256()
		s.prehasher = &p
		return nil
	case mode.Merkle:
		s.mode = m
		p := cmhash.NewSHAKE256MerkleHasher(merkle.DefaultChunkLen, cmhash.DefaultNumCPU)
		s.prehasher = &p
		return nil
	default:
		return errors.New("unsupported mode")
	}
}

func (s *Signer) Mode() mode.Mode {
	return s.mode
}

func (s *Signer) Prehasher() *hash.Hash {
	return s.prehasher
}

func (s *Signer) Algo() string {
	switch s.mode {
	case mode.Pure:
		return SLHDSASHAKE256s
	case mode.Prehash:
		return SLHDSASHAKE256sPh
	case mode.Merkle:
		return SLHDSASHAKE256sMerkle
	}
	return ""
}

func (s *Signer) Public() crypto.PublicKey {
	return s.PK
}

func (s *Signer) Sign(
	rand io.Reader,
	msg []byte,
	opts crypto.SignerOpts,
) (signature []byte, err error) {
	return circl.SignRandomized(&s.SK, rand, circl.NewMessage(msg), []byte(s.Algo()))
}

func NewSigner(v []byte) (prv *Signer, pub []byte, err error) {
	signer := Signer{}
	signer.SK.ID = circl.SHAKE_256s
	err = signer.SK.UnmarshalBinary(v)
	if err != nil {
		return
	}
	signer.PK = signer.SK.PublicKey()
	prv = &signer
	return
}
