package protocol

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"io"
	"net"

	"ananta-ps/pkg/crypto/chacha8"
)

// HandleHandshake processes client C2S handshake and performs S2C key exchange.
// It returns read and write ciphers initialized with negotiated ChaCha8 keys.
func HandleHandshake(conn net.Conn) (*chacha8.UxChaCha8, *chacha8.UxChaCha8, error) {
	// Read frame 0 (mode 0x01 plaintext)
	frame, err := ReadFrame(conn, nil)
	if err != nil {
		return nil, nil, err
	}
	if frame.Mode != 0x01 {
		return nil, nil, errors.New("invalid handshake frame mode")
	}
	if len(frame.Payload) < 20 {
		return nil, nil, errors.New("client handshake too short")
	}

	aesKey := frame.Payload[4:20] // 16-byte key

	clientWriteKey := make([]byte, 32)
	clientReadKey := make([]byte, 32)
	streamNonce := make([]byte, 12)
	aesNonce := make([]byte, 12)

	if _, err := io.ReadFull(rand.Reader, clientWriteKey); err != nil {
		return nil, nil, err
	}
	if _, err := io.ReadFull(rand.Reader, clientReadKey); err != nil {
		return nil, nil, err
	}
	if _, err := io.ReadFull(rand.Reader, streamNonce); err != nil {
		return nil, nil, err
	}
	if _, err := io.ReadFull(rand.Reader, aesNonce); err != nil {
		return nil, nil, err
	}

	chaChaHead := make([]byte, 76)
	copy(chaChaHead[0:32], clientWriteKey)
	copy(chaChaHead[32:64], clientReadKey)
	copy(chaChaHead[64:76], streamNonce)

	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, nil, err
	}
	aesGcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, err
	}

	encrypted := aesGcm.Seal(nil, aesNonce, chaChaHead, nil) // 76 + 16 = 92 bytes

	payload := make([]byte, 368)
	binary.LittleEndian.PutUint32(payload[0:4], 1)
	binary.LittleEndian.PutUint32(payload[4:8], 10)
	copy(payload[8:20], aesNonce)
	copy(payload[20:96], encrypted[0:76])
	copy(payload[96:112], encrypted[76:92])

	// Send handshake reply in plaintext with mode 0x01
	if err := WriteFrame(conn, 0x01, payload, nil); err != nil {
		return nil, nil, err
	}

	readCryptor, err := chacha8.New(clientWriteKey, streamNonce)
	if err != nil {
		return nil, nil, err
	}
	writeCryptor, err := chacha8.New(clientReadKey, streamNonce)
	if err != nil {
		return nil, nil, err
	}

	return readCryptor, writeCryptor, nil
}