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

import (
	"crypto/sha3"
	"hash"
	"io"

	"go.cypherpunks.su/keks/cm/hash/merkle"
)

const (
	SHAKE128Merkle = "shake128-merkle"
	SHAKE256Merkle = "shake256-merkle"

	SLHDSASHAKE256s       = "slh-dsa-shake-256s"
	SLHDSASHAKE256sPh     = "slh-dsa-shake-256s-ph"
	SLHDSASHAKE256sMerkle = "slh-dsa-shake-256s-merkle"
)

type SHAKE struct {
	xof *sha3.SHAKE
	l   int
}

func (h SHAKE) Size() int {
	return h.l
}

func (h SHAKE) BlockSize() int {
	return h.xof.BlockSize()
}

func (h SHAKE) Reset() {
	h.xof.Reset()
}

func (h SHAKE) Write(p []byte) (int, error) {
	return h.xof.Write(p)
}

func (h SHAKE) Sum(b []byte) []byte {
	buf := make([]byte, h.l)
	if _, err := io.ReadFull(h.xof, buf); err != nil {
		panic(err)
	}
	return append(b, buf...)
}

func NewSHAKE128() hash.Hash {
	return SHAKE{xof: sha3.NewSHAKE128(), l: 32}
}

func NewSHAKE256() hash.Hash {
	return SHAKE{xof: sha3.NewSHAKE256(), l: 64}
}

func NewCSHAKE128(s []byte) hash.Hash {
	return SHAKE{xof: sha3.NewCSHAKE128(nil, s), l: 32}
}

func NewCSHAKE256(s []byte) hash.Hash {
	return SHAKE{xof: sha3.NewCSHAKE256(nil, s), l: 64}
}

func NewSHAKE128MerkleHasher(chunkLen, workers int) hash.Hash {
	return merkle.NewHasher(
		func() hash.Hash { return NewCSHAKE128([]byte(merkle.Leaf)) },
		func() hash.Hash { return NewCSHAKE128([]byte(merkle.Node)) },
		func(h hash.Hash) hash.Hash {
			h.Reset()
			return h
		},
		func(h hash.Hash) hash.Hash {
			h.Reset()
			return h
		},
		chunkLen,
		workers,
	)
}

func NewSHAKE256MerkleHasher(chunkLen, workers int) hash.Hash {
	return merkle.NewHasher(
		func() hash.Hash { return NewCSHAKE256([]byte(merkle.Leaf)) },
		func() hash.Hash { return NewCSHAKE256([]byte(merkle.Node)) },
		func(h hash.Hash) hash.Hash {
			h.Reset()
			return h
		},
		func(h hash.Hash) hash.Hash {
			h.Reset()
			return h
		},
		chunkLen,
		workers,
	)
}
