package wire

import (
	"crypto/sha3"
	"encoding/hex"

	"golang.org/x/crypto/chacha20"
)

// Direction labels for per-direction key derivation (PROTOCOL.md §3).
const (
	DirC2S byte = 0x01 // client -> server
	DirS2C byte = 0x02 // server -> client
)

// SHA3_224 returns the SHA3-224 digest of in.
func SHA3_224(in []byte) []byte {
	h := sha3.New224()
	h.Write(in)
	return h.Sum(nil)
}

// PSKAddress is the Handshake Server Address for Intent 17: hex(SHA3-224(PSK)).
func PSKAddress(psk []byte) string {
	return hex.EncodeToString(SHA3_224(psk))
}

// DeriveKey computes the 32-byte ChaCha20 key: SHA3-256(phaseKey || dir).
func DeriveKey(phaseKey []byte, dir byte) []byte {
	h := sha3.New256()
	h.Write(phaseKey)
	h.Write([]byte{dir})
	return h.Sum(nil)
}

// CipherFor builds a ChaCha20 stream cipher for the given phase key and
// direction. ChaCha20 is symmetric, so the same cipher encrypts and decrypts.
func CipherFor(phaseKey []byte, dir byte) *chacha20.Cipher {
	c, err := chacha20.NewUnauthenticatedCipher(DeriveKey(phaseKey, dir), make([]byte, 12))
	if err != nil {
		panic("wire: chacha20 init: " + err.Error())
	}
	return c
}