package vl1

import (
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"

	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/sha3"
)

type NoiseCipher struct {
	sendCipher cipher.AEAD
	recvCipher cipher.AEAD
	sendNonce  uint64
	recvNonce  uint64
}

func NewNoiseCipher(sendKey, recvKey [32]byte) *NoiseCipher {
	send, err := chacha20poly1305.New(sendKey[:])
	if err != nil {
		panic(err)
	}
	recv, err := chacha20poly1305.New(recvKey[:])
	if err != nil {
		panic(err)
	}
	return &NoiseCipher{
		sendCipher: send,
		recvCipher: recv,
	}
}

func DeriveKeysFromPSK(psk string, localPub, remotePub []byte) ([32]byte, [32]byte) {
	h := sha3.New256()
	h.Write([]byte(psk))
	h.Write(localPub)
	h.Write(remotePub)
	sum := h.Sum(nil)

	var sendKey, recvKey [32]byte
	copy(sendKey[:], sum[:32])

	h.Reset()
	h.Write(sum)
	h.Write([]byte("reverse"))
	rev := h.Sum(nil)
	copy(recvKey[:], rev[:32])

	return sendKey, recvKey
}

func (nc *NoiseCipher) Encrypt(plaintext []byte) ([]byte, error) {
	nonce := make([]byte, 12)
	binary.BigEndian.PutUint64(nonce[4:], nc.sendNonce)
	nc.sendNonce++

	ciphertext := nc.sendCipher.Seal(nil, nonce, plaintext, nil)
	return ciphertext, nil
}

func (nc *NoiseCipher) Decrypt(ciphertext []byte) ([]byte, error) {
	nonce := make([]byte, 12)
	binary.BigEndian.PutUint64(nonce[4:], nc.recvNonce)
	nc.recvNonce++

	plaintext, err := nc.recvCipher.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt: %w", err)
	}
	return plaintext, nil
}

func (nc *NoiseCipher) EncryptTo(buf []byte, plaintext []byte) (int, error) {
	nonce := make([]byte, 12)
	binary.BigEndian.PutUint64(nonce[4:], nc.sendNonce)
	nc.sendNonce++

	ciphertext := nc.sendCipher.Seal(buf[:0], nonce, plaintext, nil)
	return len(ciphertext), nil
}

func (nc *NoiseCipher) DecryptTo(buf []byte, ciphertext []byte) ([]byte, error) {
	nonce := make([]byte, 12)
	binary.BigEndian.PutUint64(nonce[4:], nc.recvNonce)
	nc.recvNonce++

	plaintext, err := nc.recvCipher.Open(buf[:0], nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt: %w", err)
	}
	return plaintext, nil
}

func generateSessionKey() [32]byte {
	var key [32]byte
	if _, err := io.ReadFull(rand.Reader, key[:]); err != nil {
		panic(err)
	}
	return key
}