package cql

import (
	"bytes"
	"testing"

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

func TestCompressedSegmentReadWrite(t *testing.T) {
	payload := bytes.Repeat([]byte{0xAB, 0xCD}, 64)
	segment := &Segment{
		Payload:       payload,
		SelfContained: true,
		Compressed:    true,
	}

	var buf bytes.Buffer
	compressor := &lz4.LZ4Compressor{}
	if err := segment.WriteSegment(&buf, compressor); err != nil {
		t.Fatalf("WriteSegment failed: %v", err)
	}

	readSegment, err := ReadSegment(&buf, compressor)
	if err != nil {
		t.Fatalf("ReadSegment failed: %v", err)
	}

	if !bytes.Equal(readSegment.Payload, payload) {
		t.Fatalf("Payload mismatch: got %x want %x", readSegment.Payload, payload)
	}
	if !readSegment.Compressed {
		t.Fatalf("Expected compressed segment")
	}
}

func TestSegmentReaderMultiFrameBuffer(t *testing.T) {
	frame1 := &Frame{
		Version: ProtoVersion5,
		Flags:   0,
		Stream:  1,
		Opcode:  OpcodeReady,
		Body:    []byte("one"),
		Length:  int32(len("one")),
	}
	frame2 := &Frame{
		Version: ProtoVersion5,
		Flags:   0,
		Stream:  2,
		Opcode:  OpcodeReady,
		Body:    []byte("two"),
		Length:  int32(len("two")),
	}

	var frameBuf bytes.Buffer
	if err := frame1.WriteFrame(&frameBuf); err != nil {
		t.Fatalf("WriteFrame failed: %v", err)
	}
	if err := frame2.WriteFrame(&frameBuf); err != nil {
		t.Fatalf("WriteFrame failed: %v", err)
	}

	segment := &Segment{
		Payload:       frameBuf.Bytes(),
		SelfContained: true,
		Compressed:    false,
	}

	var segmentBuf bytes.Buffer
	if err := segment.WriteSegment(&segmentBuf, nil); err != nil {
		t.Fatalf("WriteSegment failed: %v", err)
	}

	reader := NewSegmentReader(&segmentBuf, nil)
	first, err := reader.ReadFrame()
	if err != nil {
		t.Fatalf("ReadFrame failed: %v", err)
	}
	second, err := reader.ReadFrame()
	if err != nil {
		t.Fatalf("ReadFrame failed: %v", err)
	}

	firstFrame, err := ParseFrame(first)
	if err != nil {
		t.Fatalf("ParseFrame failed: %v", err)
	}
	secondFrame, err := ParseFrame(second)
	if err != nil {
		t.Fatalf("ParseFrame failed: %v", err)
	}

	if string(firstFrame.Body) != "one" {
		t.Fatalf("First frame body mismatch: got %q", string(firstFrame.Body))
	}
	if string(secondFrame.Body) != "two" {
		t.Fatalf("Second frame body mismatch: got %q", string(secondFrame.Body))
	}
}

func TestSegmentWriterLargeFrameUsesNonSelfContainedSlices(t *testing.T) {
	// Build a frame larger than one segment so writer must split it.
	body := bytes.Repeat([]byte{0x5A}, MaxSegmentPayload+256)
	frame := &Frame{
		Version: ProtoVersion5,
		Flags:   0,
		Stream:  7,
		Opcode:  OpcodeResult,
		Body:    body,
		Length:  int32(len(body)),
	}

	var frameBuf bytes.Buffer
	if err := frame.WriteFrame(&frameBuf); err != nil {
		t.Fatalf("WriteFrame failed: %v", err)
	}
	rawFrame := frameBuf.Bytes()

	var segBuf bytes.Buffer
	writer := NewSegmentWriter(&segBuf, nil)
	if err := writer.WriteFrame(rawFrame, false); err != nil {
		t.Fatalf("WriteFrame (segment writer) failed: %v", err)
	}

	// Read back each segment and ensure all slices are marked non-self-contained,
	// while concatenated payload reassembles the original frame bytes.
	var reassembled bytes.Buffer
	segments := 0
	for segBuf.Len() > 0 {
		seg, err := ReadSegment(&segBuf, nil)
		if err != nil {
			t.Fatalf("ReadSegment failed: %v", err)
		}
		segments++
		if seg.SelfContained {
			t.Fatalf("fragmented segment %d should not be self-contained", segments)
		}
		reassembled.Write(seg.Payload)
	}
	if segments < 2 {
		t.Fatalf("expected multiple segments, got %d", segments)
	}
	if !bytes.Equal(reassembled.Bytes(), rawFrame) {
		t.Fatalf("reassembled frame mismatch: got %d bytes want %d", reassembled.Len(), len(rawFrame))
	}
}