package owner

// 5% of AI generated code
// human made; AI details
// FINAL-FINAL

import (
	"WhspBrd/typio/bit"
	"WhspBrd/typio/yum"
	"bytes"
	"crypto/mlkem"
	"crypto/sha256"
	"errors"

	"github.com/AeonDave/cryptonite-go/sig"
)

var mldsa87 = sig.NewDeterministicMLDSA87()

var (
	ErrInvalidSignature = errors.New("invalid signature")
	ErrDestroyedSecret  = errors.New("secret has been destroyed")
)

const requiredSeedSizeForMLDSA87 = 32

const signaturePubPartSize = sig.MLDSA87PublicKeySize
const signatureSignSize = sig.MLDSA87SignatureSize
const SignatureSize = signaturePubPartSize + signatureSignSize

const IdentitySize = bit.Size256b_B

type Identity = bit.Sha256

func IdentityEq(id, other Identity) bool {
	return bytes.Equal(id[:], other[:])
}

func pubHash(pubKey []byte) Identity {
	return sha256.Sum256(pubKey)
}

type Secret interface {
	Identity() Identity
	PulicKey() []byte
	EncapKey() []byte
	Decapsulate(ciphertext []byte) ([]byte, error)
	Sign(data []byte) (signedData []byte, err error)
	Save(path string, password []byte) error
	Destroy()
	SaveAndDestroy(path string, password []byte) error
}

type secret struct {
	seed      []byte
	pubKey    []byte
	encap     []byte
	identity  Identity
	decap     *mlkem.DecapsulationKey1024
	privKey   []byte
	destroyed bool
}

func Encapsulate(secret Secret, encapKey []byte) (sharedkey []byte, ciphertext []byte, err error) {
	encap, err := mlkem.NewEncapsulationKey1024(encapKey)
	if err != nil {
		return
	}
	sharedkey, ciphertext = encap.Encapsulate()
	return
}

func NewSecret() (Secret, error) {
	seed := make([]byte, requiredSeedSizeForMLDSA87)
	if err := yum.YumSeed(seed); err != nil {
		return nil, err
	}
	pubKey, privKey, err := sig.GenerateDeterministicKeyMLDSA87(seed)
	if err != nil {
		return nil, err
	}
	decap, err := mlkem.NewDecapsulationKey1024(append(seed, seed...))
	if err != nil {
		return nil, err
	}
	encap := decap.EncapsulationKey().Bytes()
	return &secret{
		seed:     seed,
		encap:    encap,
		decap:    decap,
		pubKey:   pubKey,
		identity: pubHash(pubKey),
		privKey:  privKey,
	}, nil
}

func LoadSecret(path string, password []byte) (Secret, error) {
	seed, err := yum.YumLoad(path, password)
	if err != nil {
		return nil, err
	}
	pubKey, privKey, err := sig.GenerateDeterministicKeyMLDSA87(seed)
	if err != nil {
		return nil, err
	}
	decap, err := mlkem.NewDecapsulationKey1024(append(seed, seed...))
	if err != nil {
		return nil, err
	}
	encap := decap.EncapsulationKey().Bytes()
	return &secret{
		seed:     seed,
		encap:    encap,
		decap:    decap,
		pubKey:   pubKey,
		identity: pubHash(pubKey),
		privKey:  privKey,
	}, nil
}

func (s *secret) Identity() Identity {
	return s.identity
}

func (s *secret) PulicKey() []byte {
	pubCopy := make([]byte, len(s.pubKey))
	copy(pubCopy, s.pubKey)
	return pubCopy
}

func (s *secret) EncapKey() []byte {
	encapCopy := make([]byte, len(s.encap))
	copy(encapCopy, s.encap)
	return encapCopy
}

func (s *secret) Decapsulate(ciphertext []byte) ([]byte, error) {
	if s.destroyed {
		return nil, ErrDestroyedSecret
	}
	return s.decap.Decapsulate(ciphertext)
}

func (s *secret) Sign(data []byte) (sign []byte, err error) {
	if s.destroyed {
		err = ErrDestroyedSecret
		return
	}

	signPart, err := mldsa87.Sign(s.privKey, data)
	if err != nil {
		return
	}

	sign = append(s.pubKey, signPart...)
	return
}

func (s *secret) Save(path string, password []byte) error {
	if s.destroyed {
		return ErrDestroyedSecret
	}
	return yum.YumSave(path, s.seed, password)
}

func Verify(data []byte, sign []byte) (Identity, error) {
	if len(sign) < SignatureSize {
		return Identity{}, ErrInvalidSignature
	}

	pubPart := sign[:signaturePubPartSize]
	signPart := sign[signaturePubPartSize:]

	ok := mldsa87.Verify(pubPart, data, signPart)
	if !ok {
		return Identity{}, ErrInvalidSignature
	}

	return pubHash(pubPart), nil
}

// Destroy zeroes out the secret's sensitive data and marks it as destroyed.
func (s *secret) Destroy() {
	if s.destroyed {
		return
	}
	for i := range s.seed {
		s.seed[i] = 0
	}
	for i := range s.privKey {
		s.privKey[i] = 0
	}
	for i := range s.encap {
		s.encap[i] = 0
	}
	s.destroyed = true
	/*for i := range s.pubKey {
		s.pubKey[i] = 0
	}
	for i := range s.identity {
		s.identity[i] = 0
	}*/ // not needed for added security, but could be done if desired
}

func (s *secret) SaveAndDestroy(path string, password []byte) error {
	if s.destroyed {
		return ErrDestroyedSecret
	}
	if err := s.Save(path, password); err != nil {
		return err
	}
	s.Destroy()
	return nil
}