package eq2net

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"io"
)

const (
	// Compression flags
	CompressionFlagZlib   = 0x5A // Zlib compression
	CompressionFlagSimple = 0xA5 // Simple encoding (no actual compression)
	
	// Compression threshold - packets larger than this use zlib
	CompressionThreshold = 30
)

// CompressPacket compresses a packet using zlib or simple encoding
func CompressPacket(data []byte) ([]byte, error) {
	if len(data) < 2 {
		return data, nil
	}
	
	// Determine opcode size
	flagOffset := 1
	if data[0] == 0 {
		flagOffset = 2 // Two-byte opcode
	}
	
	// Don't compress if too small
	if len(data) <= flagOffset {
		return data, nil
	}
	
	result := make([]byte, 0, len(data)+1)
	
	// Copy opcode bytes
	result = append(result, data[:flagOffset]...)
	
	if len(data) > CompressionThreshold {
		// Use zlib compression for larger packets
		result = append(result, CompressionFlagZlib)
		
		// Compress the data after opcode
		var compressed bytes.Buffer
		w := zlib.NewWriter(&compressed)
		if _, err := w.Write(data[flagOffset:]); err != nil {
			return nil, err
		}
		if err := w.Close(); err != nil {
			return nil, err
		}
		
		result = append(result, compressed.Bytes()...)
	} else {
		// Use simple encoding for smaller packets
		result = append(result, CompressionFlagSimple)
		result = append(result, data[flagOffset:]...)
	}
	
	return result, nil
}

// DecompressPacket decompresses a packet
func DecompressPacket(data []byte) ([]byte, error) {
	if len(data) < 3 {
		return data, nil
	}
	
	// Determine opcode size and compression flag position
	flagOffset := 1
	if data[0] == 0 {
		flagOffset = 2
	}
	
	if len(data) <= flagOffset {
		return data, nil
	}
	
	compressionFlag := data[flagOffset]
	
	// Check compression type
	switch compressionFlag {
	case CompressionFlagZlib:
		// Zlib decompression
		result := make([]byte, 0, len(data)*2)
		
		// Copy opcode
		result = append(result, data[:flagOffset]...)
		
		// Decompress data (skip flag byte)
		r, err := zlib.NewReader(bytes.NewReader(data[flagOffset+1:]))
		if err != nil {
			return nil, err
		}
		defer r.Close()
		
		decompressed, err := io.ReadAll(r)
		if err != nil {
			return nil, err
		}
		
		result = append(result, decompressed...)
		return result, nil
		
	case CompressionFlagSimple:
		// Simple encoding - just remove the flag byte
		result := make([]byte, 0, len(data)-1)
		result = append(result, data[:flagOffset]...)
		result = append(result, data[flagOffset+1:]...)
		return result, nil
		
	default:
		// No compression
		return data, nil
	}
}

// IsCompressed checks if a packet is compressed
func IsCompressed(data []byte) bool {
	if len(data) < 2 {
		return false
	}
	
	flagOffset := 1
	if data[0] == 0 {
		flagOffset = 2
	}
	
	if len(data) <= flagOffset {
		return false
	}
	
	flag := data[flagOffset]
	return flag == CompressionFlagZlib || flag == CompressionFlagSimple
}

// ChatEncode encodes chat data using XOR encryption with rolling key
func ChatEncode(data []byte, encodeKey uint32) []byte {
	// Skip certain packet types
	if len(data) >= 2 && (data[1] == 0x01 || data[0] == 0x02 || data[0] == 0x1d) {
		return data
	}
	
	// Work with data after opcode
	if len(data) <= 2 {
		return data
	}
	
	result := make([]byte, len(data))
	copy(result[:2], data[:2]) // Copy opcode
	
	key := encodeKey
	offset := 2
	
	// Process 4-byte blocks with rolling key
	for i := offset; i+4 <= len(data); i += 4 {
		block := binary.LittleEndian.Uint32(data[i : i+4])
		encrypted := block ^ key
		binary.LittleEndian.PutUint32(result[i:i+4], encrypted)
		key = encrypted // Update key with encrypted data
	}
	
	// Handle remaining bytes
	keyByte := byte(key & 0xFF)
	alignedEnd := offset + ((len(data)-offset)/4)*4
	for i := alignedEnd; i < len(data); i++ {
		result[i] = data[i] ^ keyByte
	}
	
	return result
}

// ChatDecode decodes chat data using XOR encryption with rolling key
func ChatDecode(data []byte, decodeKey uint32) []byte {
	// Skip certain packet types
	if len(data) >= 2 && (data[1] == 0x01 || data[0] == 0x02 || data[0] == 0x1d) {
		return data
	}
	
	// Work with data after opcode
	if len(data) <= 2 {
		return data
	}
	
	result := make([]byte, len(data))
	copy(result[:2], data[:2]) // Copy opcode
	
	key := decodeKey
	offset := 2
	
	// Process 4-byte blocks with rolling key
	for i := offset; i+4 <= len(data); i += 4 {
		encrypted := binary.LittleEndian.Uint32(data[i : i+4])
		decrypted := encrypted ^ key
		binary.LittleEndian.PutUint32(result[i:i+4], decrypted)
		key = encrypted // Update key with encrypted data (before decryption)
	}
	
	// Handle remaining bytes
	keyByte := byte(key & 0xFF)
	alignedEnd := offset + ((len(data)-offset)/4)*4
	for i := alignedEnd; i < len(data); i++ {
		result[i] = data[i] ^ keyByte
	}
	
	return result
}