package udp

import (
	"crypto/rc4"
	"fmt"
)

// 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{
		encrypted: false,
	}
}

// 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 fmt.Errorf("encryption key cannot be empty")
	}

	// Create separate RC4 ciphers for bidirectional communication
	clientCipher, err := rc4.NewCipher(key)
	if err != nil {
		return fmt.Errorf("failed to create client cipher: %w", err)
	}

	serverCipher, err := rc4.NewCipher(key)
	if err != nil {
		return fmt.Errorf("failed to create server cipher: %w", 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
}

// Clone creates a copy of the crypto instance with the same key
func (c *Crypto) Clone() (*Crypto, error) {
	newCrypto := NewCrypto()
	if c.encrypted && c.key != nil {
		return newCrypto, newCrypto.SetKey(c.key)
	}
	return newCrypto, nil
}