package packets

import (
	"crypto/rand"
	"crypto/rc4"
	"crypto/rsa"
	"crypto/x509"
	"encoding/binary"
	"fmt"

	"eq2emu/internal/crc16"
	"eq2emu/internal/opcodes"
)

type SecurityManager struct {
	sessionKey uint32
	rc4Cipher  *rc4.Cipher
	rsaKey     *rsa.PrivateKey
	encrypted  bool
	compressed bool
}

func NewSecurityManager() *SecurityManager {
	// Generate RSA key for session establishment
	key, _ := rsa.GenerateKey(rand.Reader, 1024)

	return &SecurityManager{
		rsaKey:    key,
		encrypted: false,
	}
}

// Validate packet CRC
func (sm *SecurityManager) ValidateCRC(data []byte, key uint32) bool {
	if len(data) < 2 {
		return false
	}

	// Skip CRC validation for session packets
	if len(data) >= 2 {
		opcode := opcodes.ProtocolOpcode(data[0])
		if opcode == opcodes.OP_SessionRequest ||
			opcode == opcodes.OP_SessionResponse ||
			opcode == opcodes.OP_OutOfSession {
			return true
		}
	}

	// Extract CRC from last 2 bytes
	packetCRC := binary.LittleEndian.Uint16(data[len(data)-2:])
	if packetCRC == 0 {
		return true // No CRC
	}

	// Calculate CRC on data without the CRC bytes
	calcCRC := uint16(crc16.Calculate(data[:len(data)-2], key))
	return calcCRC == packetCRC
}

// Append CRC to packet
func (sm *SecurityManager) AppendCRC(data []byte, key uint32) []byte {
	// Skip CRC for session packets
	if len(data) >= 1 {
		opcode := opcodes.ProtocolOpcode(data[0])
		if opcode == opcodes.OP_SessionRequest ||
			opcode == opcodes.OP_SessionResponse ||
			opcode == opcodes.OP_OutOfSession {
			return data
		}
	}

	crc := uint16(crc16.Calculate(data, key))
	result := make([]byte, len(data)+2)
	copy(result, data)
	binary.LittleEndian.PutUint16(result[len(data):], crc)
	return result
}

// Chat encoding (custom XOR encryption)
func (sm *SecurityManager) ChatEncode(data []byte, key uint32) {
	if len(data) < 2 {
		return
	}

	// Skip encoding for certain opcodes
	if data[1] == 0x01 || data[0] == 0x02 || data[0] == 0x1d {
		return
	}

	// Skip opcode bytes
	payload := data[2:]
	encKey := key

	// Encode 4-byte chunks
	for i := 0; i+3 < len(payload); i += 4 {
		chunk := binary.LittleEndian.Uint32(payload[i:])
		encoded := chunk ^ encKey
		encKey = encoded
		binary.LittleEndian.PutUint32(payload[i:], encoded)
	}

	// Encode remaining bytes
	keyByte := uint8(encKey & 0xFF)
	for i := (len(payload) / 4) * 4; i < len(payload); i++ {
		payload[i] ^= keyByte
	}
}

// Chat decoding
func (sm *SecurityManager) ChatDecode(data []byte, key uint32) {
	if len(data) < 2 {
		return
	}

	// Skip decoding for certain opcodes
	if data[1] == 0x01 || data[0] == 0x02 || data[0] == 0x1d {
		return
	}

	// Skip opcode bytes
	payload := data[2:]
	decKey := key

	// Decode 4-byte chunks
	for i := 0; i+3 < len(payload); i += 4 {
		chunk := binary.LittleEndian.Uint32(payload[i:])
		decoded := chunk ^ decKey
		decKey = chunk
		binary.LittleEndian.PutUint32(payload[i:], decoded)
	}

	// Decode remaining bytes
	keyByte := uint8(decKey & 0xFF)
	for i := (len(payload) / 4) * 4; i < len(payload); i++ {
		payload[i] ^= keyByte
	}
}

// RC4 encryption setup
func (sm *SecurityManager) SetRC4Key(key []byte) error {
	cipher, err := rc4.NewCipher(key)
	if err != nil {
		return err
	}
	sm.rc4Cipher = cipher
	sm.encrypted = true
	return nil
}

// RC4 encrypt/decrypt
func (sm *SecurityManager) RC4Crypt(data []byte) {
	if sm.rc4Cipher != nil {
		sm.rc4Cipher.XORKeyStream(data, data)
	}
}

// RSA decrypt (for key exchange)
func (sm *SecurityManager) RSADecrypt(data []byte) ([]byte, error) {
	if sm.rsaKey == nil {
		return nil, fmt.Errorf("no RSA key")
	}

	return rsa.DecryptPKCS1v15(rand.Reader, sm.rsaKey, data)
}

// Get RSA public key for client
func (sm *SecurityManager) GetRSAPublicKey() ([]byte, error) {
	if sm.rsaKey == nil {
		return nil, fmt.Errorf("no RSA key")
	}

	return x509.MarshalPKIXPublicKey(&sm.rsaKey.PublicKey)
}

// Process encrypted packet
func (sm *SecurityManager) ProcessEncryptedPacket(data []byte) []byte {
	if len(data) < 3 {
		return data
	}

	// Make copy to avoid modifying original
	result := make([]byte, len(data))
	copy(result, data)

	// Apply RC4 if enabled
	if sm.encrypted && sm.rc4Cipher != nil {
		sm.RC4Crypt(result)
	}

	// Apply chat decoding
	if sm.sessionKey != 0 {
		sm.ChatDecode(result, sm.sessionKey)
	}

	return result
}

// Prepare packet for sending
func (sm *SecurityManager) PrepareOutgoingPacket(data []byte) []byte {
	// Make copy
	result := make([]byte, len(data))
	copy(result, data)

	// Apply chat encoding
	if sm.sessionKey != 0 {
		sm.ChatEncode(result, sm.sessionKey)
	}

	// Apply RC4 if enabled
	if sm.encrypted && sm.rc4Cipher != nil {
		sm.RC4Crypt(result)
	}

	// Append CRC
	result = sm.AppendCRC(result, sm.sessionKey)

	return result
}

// Session management
func (sm *SecurityManager) SetSessionKey(key uint32) {
	sm.sessionKey = key
}

func (sm *SecurityManager) GetSessionKey() uint32 {
	return sm.sessionKey
}

func (sm *SecurityManager) SetEncrypted(encrypted bool) {
	sm.encrypted = encrypted
}

func (sm *SecurityManager) IsEncrypted() bool {
	return sm.encrypted
}

func (sm *SecurityManager) SetCompressed(compressed bool) {
	sm.compressed = compressed
}

func (sm *SecurityManager) IsCompressed() bool {
	return sm.compressed
}