package udp

import (
	"crypto/rc4"
	"errors"
)

// EQ2EMu CRC32 polynomial (reversed)
const crcPolynomial = 0xEDB88320

// Pre-computed CRC32 lookup table for fast calculation
var crcTable [256]uint32

// init builds the CRC lookup table at package initialization
func init() {
	for i := range crcTable {
		crc := uint32(i)
		for range 8 {
			if crc&1 == 1 {
				crc = (crc >> 1) ^ crcPolynomial
			} else {
				crc >>= 1
			}
		}
		crcTable[i] = crc
	}
}

// CalculateCRC32 computes CRC32 using EQ2EMu's algorithm
// Returns 16-bit value by truncating the upper bits
func CalculateCRC32(data []byte) uint16 {
	crc := uint32(0xFFFFFFFF)

	// Use lookup table for efficient calculation
	for _, b := range data {
		crc = crcTable[byte(crc)^b] ^ (crc >> 8)
	}

	// Return inverted result truncated to 16 bits
	return uint16(^crc)
}

// ValidateCRC checks if packet has valid CRC
// Expects CRC to be the last 2 bytes of data
func ValidateCRC(data []byte) bool {
	if len(data) < 2 {
		return false
	}

	// Split payload and CRC
	payload := data[:len(data)-2]
	expectedCRC := uint16(data[len(data)-2]) | (uint16(data[len(data)-1]) << 8)

	// Calculate and compare
	actualCRC := CalculateCRC32(payload)
	return expectedCRC == actualCRC
}

// AppendCRC adds 16-bit CRC to the end of data
func AppendCRC(data []byte) []byte {
	crc := CalculateCRC32(data)
	result := make([]byte, len(data)+2)
	copy(result, data)

	// Append CRC in little-endian format
	result[len(data)] = byte(crc)
	result[len(data)+1] = byte(crc >> 8)

	return result
}

// ValidateAndStrip validates CRC and returns data without CRC suffix
func ValidateAndStrip(data []byte) ([]byte, bool) {
	if !ValidateCRC(data) {
		return nil, false
	}
	return data[:len(data)-2], true
}

// Crypto handles RC4 encryption/decryption for EQ2EMu protocol
type Crypto struct {
	clientCipher *rc4.Cipher // Cipher for decrypting client data
	serverCipher *rc4.Cipher // Cipher for encrypting server data
	key          []byte      // Encryption key
	encrypted    bool        // Whether encryption is active
}

// NewCrypto creates a new crypto instance with encryption disabled
func NewCrypto() *Crypto {
	return &Crypto{}
}

// SetKey initializes RC4 encryption with the given key
// Creates separate ciphers for client and server with 20-byte priming
func (c *Crypto) SetKey(key []byte) error {
	if len(key) == 0 {
		return errors.New("encryption key cannot be empty")
	}

	// Create separate RC4 ciphers for bidirectional communication
	clientCipher, err := rc4.NewCipher(key)
	if err != nil {
		return err
	}

	serverCipher, err := rc4.NewCipher(key)
	if err != nil {
		return err
	}

	// Prime both ciphers with 20 dummy bytes per EQ2EMu protocol
	dummy := make([]byte, 20)
	clientCipher.XORKeyStream(dummy, dummy)
	serverCipher.XORKeyStream(dummy, dummy)

	c.clientCipher = clientCipher
	c.serverCipher = serverCipher
	c.key = make([]byte, len(key))
	copy(c.key, key)
	c.encrypted = true

	return nil
}

// IsEncrypted returns whether encryption is currently active
func (c *Crypto) IsEncrypted() bool {
	return c.encrypted
}

// Encrypt encrypts data for transmission to client
func (c *Crypto) Encrypt(data []byte) []byte {
	if !c.encrypted || c.serverCipher == nil {
		return data
	}

	encrypted := make([]byte, len(data))
	copy(encrypted, data)
	c.serverCipher.XORKeyStream(encrypted, encrypted)
	return encrypted
}

// Decrypt decrypts data received from client
func (c *Crypto) Decrypt(data []byte) []byte {
	if !c.encrypted || c.clientCipher == nil {
		return data
	}

	decrypted := make([]byte, len(data))
	copy(decrypted, data)
	c.clientCipher.XORKeyStream(decrypted, decrypted)
	return decrypted
}

// GetKey returns a copy of the encryption key
func (c *Crypto) GetKey() []byte {
	if c.key == nil {
		return nil
	}
	keyCopy := make([]byte, len(c.key))
	copy(keyCopy, c.key)
	return keyCopy
}

// Reset disables encryption and clears keys
func (c *Crypto) Reset() {
	c.clientCipher = nil
	c.serverCipher = nil
	c.key = nil
	c.encrypted = false
}
