package cookie

import (
	"crypto/hmac"
	"crypto/sha256"
	"errors"
	"hash"
)

type KeyOrPassword struct {
	SecretType  int
	Key         *Key
	Password    string
	DerivedKeys *DerivedKeys
	Akey        []byte
	Ekey        []byte
}

type DerivedKeys struct{ akey, ekey []byte }

type Key struct{ bytes []byte }

func (k *Key) GetRawBytes() []byte { return k.bytes }

func (kp *KeyOrPassword) DeriveKeys(salt []byte) (*DerivedKeys, error) {
	if len(salt) != SaltSize {
		return nil, errors.New("bad salt")
	}
	if kp.SecretType == 1 {
		akey := kp.hkdf(sha256.New, kp.Key.GetRawBytes(), KeyByteSize, AuthInfo, salt)
		ekey := kp.hkdf(sha256.New, kp.Key.GetRawBytes(), KeyByteSize, EncInfo, salt)
		return &DerivedKeys{akey: akey, ekey: ekey}, nil
	}
	return nil, errors.New("unsupported")
}

// === HKDF ===
func (kp *KeyOrPassword) hkdf(hashFunc func() hash.Hash, ikm []byte, length int, info string, salt []byte) []byte {
	digestLen := hashFunc().Size()
	if salt == nil {
		salt = make([]byte, digestLen)
	}

	// Extract
	prkMac := hmac.New(hashFunc, salt)
	prkMac.Write(ikm)
	prk := prkMac.Sum(nil)

	// Expand
	var okm []byte
	prev := []byte{}
	counter := byte(1)

	for len(okm) < length {
		h := hmac.New(hashFunc, prk)
		h.Write(prev)
		h.Write([]byte(info))
		h.Write([]byte{counter})
		step := h.Sum(nil)

		okm = append(okm, step...)
		prev = step
		counter++
	}

	return okm[:length]
}