package cql

import (
	"bytes"
	"strings"
	"testing"

	"github.com/apache/cassandra-gocql-driver/v2/lz4"
)

func TestMaybeDecompressFrameAllowsAuthResponseBeforeStartupComplete(t *testing.T) {
	compressor := &lz4.LZ4Compressor{}
	original := []byte("auth-bytes")
	compressed, err := compressor.AppendCompressedWithLength(nil, original)
	if err != nil {
		t.Fatalf("compress failed: %v", err)
	}

	conn := &Conn{}
	conn.stateMu.Lock()
	conn.protocolVersion = ProtoVersion4
	conn.compressionEnabled = true
	conn.compressor = compressor
	conn.startupComplete = false
	conn.stateMu.Unlock()

	frame := &Frame{
		Opcode: OpcodeAuthResponse,
		Flags:  FlagCompression,
		Body:   compressed,
		Length: int32(len(compressed)),
	}

	if err := conn.maybeDecompressFrame(frame); err != nil {
		t.Fatalf("maybeDecompressFrame failed: %v", err)
	}
	if frame.Flags&FlagCompression != 0 {
		t.Fatalf("expected compression flag to be cleared")
	}
	if !bytes.Equal(frame.Body, original) {
		t.Fatalf("unexpected decompressed body: got %q want %q", frame.Body, original)
	}
}

func TestMaybeDecompressFrameRejectsCompressedQueryBeforeStartupComplete(t *testing.T) {
	compressor := &lz4.LZ4Compressor{}
	original := []byte("query-bytes")
	compressed, err := compressor.AppendCompressedWithLength(nil, original)
	if err != nil {
		t.Fatalf("compress failed: %v", err)
	}

	conn := &Conn{}
	conn.stateMu.Lock()
	conn.protocolVersion = ProtoVersion4
	conn.compressionEnabled = true
	conn.compressor = compressor
	conn.startupComplete = false
	conn.stateMu.Unlock()

	frame := &Frame{
		Opcode: OpcodeQuery,
		Flags:  FlagCompression,
		Body:   compressed,
		Length: int32(len(compressed)),
	}

	err = conn.maybeDecompressFrame(frame)
	if err == nil {
		t.Fatalf("expected error for compressed query before startup complete")
	}
	if !strings.Contains(err.Error(), "before STARTUP completed") {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestMaybeDecompressFrameIgnoresCompressionFlagForV5(t *testing.T) {
	conn := &Conn{}
	conn.stateMu.Lock()
	conn.protocolVersion = ProtoVersion5
	conn.compressionEnabled = false
	conn.compressor = nil
	conn.startupComplete = false
	conn.stateMu.Unlock()

	frame := &Frame{
		Opcode: OpcodeQuery,
		Flags:  FlagCompression,
		Body:   []byte("raw"),
		Length: 3,
	}

	if err := conn.maybeDecompressFrame(frame); err != nil {
		t.Fatalf("maybeDecompressFrame failed: %v", err)
	}
	if frame.Flags&FlagCompression != 0 {
		t.Fatalf("expected compression flag to be cleared")
	}
}