package cql

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"

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

// Segment represents a protocol v5+ segment
type Segment struct {
	// Payload contains the frame data (may be partial or multiple frames)
	Payload []byte
	// SelfContained indicates this segment contains complete frame(s)
	SelfContained bool
	// Compressed indicates the payload is compressed
	Compressed bool
	// UncompressedLength is the original length before compression (0 if not compressed)
	UncompressedLength int
}

// ReadSegment reads a segment from the reader
func ReadSegment(r io.Reader, compressor gocql.Compressor) (*Segment, error) {
	// Read first 6 bytes to attempt uncompressed header parsing.
	header := make([]byte, 6)
	if _, err := io.ReadFull(r, header); err != nil {
		return nil, fmt.Errorf("failed to read segment header: %w", err)
	}

	// Try to parse as uncompressed first
	// Uncompressed: [PayloadLen(17 bits) | S(1 bit) | CRC24(24 bits)]
	// Byte 0: PayloadLen bits 0-7
	// Byte 1: PayloadLen bits 8-15
	// Byte 2: PayloadLen bit 16 (bit 0), S flag (bit 1)
	payloadLen := int(header[0]) | int(header[1])<<8 | int(header[2]&0x01)<<16
	selfContained := (header[2] & 0x02) != 0
	headerCRC24 := uint32(header[3]) | uint32(header[4])<<8 | uint32(header[5])<<16

	// Verify CRC24 of first 3 bytes
	computedCRC24 := Crc24(header[:3])
	if computedCRC24 == headerCRC24 {
		// This is an uncompressed segment
		payload := make([]byte, payloadLen)
		if _, err := io.ReadFull(r, payload); err != nil {
			return nil, fmt.Errorf("failed to read uncompressed payload: %w", err)
		}

		// Read and verify CRC32
		crc32Bytes := make([]byte, 4)
		if _, err := io.ReadFull(r, crc32Bytes); err != nil {
			return nil, fmt.Errorf("failed to read CRC32: %w", err)
		}
		readCRC32 := binary.LittleEndian.Uint32(crc32Bytes)
		computedCRC32 := Crc32(payload)
		if computedCRC32 != readCRC32 {
			return nil, fmt.Errorf("CRC32 mismatch: expected %d, got %d", computedCRC32, readCRC32)
		}

		return &Segment{
			Payload:            payload,
			SelfContained:      selfContained,
			Compressed:         false,
			UncompressedLength: 0,
		}, nil
	}

	// This must be a compressed segment
	extra := make([]byte, 2)
	if _, err := io.ReadFull(r, extra); err != nil {
		return nil, fmt.Errorf("failed to read segment header: %w", err)
	}
	fullHeader := append(header, extra...)
	// Compressed: [CompressedLen(17) | UncompressedLen(17) | S(1) | CRC24(24)]
	// Byte 0-1: CompressedLen bits 0-15
	// Byte 2: CompressedLen bit 16 (bit 0), UncompressedLen bits 0-6 (bits 1-7)
	// Byte 3: UncompressedLen bits 7-14
	// Byte 4: UncompressedLen bit 15 (bit 0), S flag (bit 1)
	// Bytes 5-7: CRC24
	compressedLen := int(fullHeader[0]) | int(fullHeader[1])<<8 | int(fullHeader[2]&0x01)<<16
	uncompressedLen := int(fullHeader[2]>>1) | int(fullHeader[3])<<7 | int(fullHeader[4]&0x03)<<15
	selfContained = (fullHeader[4] & 0x04) != 0
	headerCRC24 = uint32(fullHeader[5]) | uint32(fullHeader[6])<<8 | uint32(fullHeader[7])<<16

	// Verify CRC24 of first 5 bytes
	computedCRC24 = Crc24(fullHeader[:5])
	if computedCRC24 != headerCRC24 {
		return nil, fmt.Errorf("CRC24 mismatch in compressed header: expected %d, got %d", computedCRC24, headerCRC24)
	}

	// Read compressed payload
	compressedPayload := make([]byte, compressedLen)
	if _, err := io.ReadFull(r, compressedPayload); err != nil {
		return nil, fmt.Errorf("failed to read compressed payload: %w", err)
	}

	// Read and verify CRC32
	crc32Bytes := make([]byte, 4)
	if _, err := io.ReadFull(r, crc32Bytes); err != nil {
		return nil, fmt.Errorf("failed to read CRC32: %w", err)
	}
	readCRC32 := binary.LittleEndian.Uint32(crc32Bytes)
	computedCRC32 := Crc32(compressedPayload)
	if computedCRC32 != readCRC32 {
		return nil, fmt.Errorf("CRC32 mismatch: expected %d, got %d", computedCRC32, readCRC32)
	}

	// Decompress if needed
	var payload []byte
	var err error
	if uncompressedLen > 0 && compressor != nil {
		payload, err = compressor.AppendDecompressed(nil, compressedPayload, uint32(uncompressedLen))
		if err != nil {
			return nil, fmt.Errorf("failed to decompress payload: %w", err)
		}
		if len(payload) != uncompressedLen {
			return nil, fmt.Errorf("decompressed length mismatch: expected %d, got %d", uncompressedLen, len(payload))
		}
	} else {
		// No compression or uncompressed_len = 0 means use as-is
		payload = compressedPayload
		uncompressedLen = len(payload)
	}

	return &Segment{
		Payload:            payload,
		SelfContained:      selfContained,
		Compressed:         true,
		UncompressedLength: uncompressedLen,
	}, nil
}

// WriteSegment writes a segment to the writer
func (s *Segment) WriteSegment(w io.Writer, compressor gocql.Compressor) error {
	if s.Compressed {
		return s.writeCompressedSegment(w, compressor)
	}
	return s.writeUncompressedSegment(w)
}

func (s *Segment) writeUncompressedSegment(w io.Writer) error {
	payloadLen := len(s.Payload)
	if payloadLen > MaxSegmentPayload {
		return fmt.Errorf("payload too large: %d bytes (max %d)", payloadLen, MaxSegmentPayload)
	}

	// Build header
	var header [6]byte
	headerInt := uint32(payloadLen)
	if s.SelfContained {
		headerInt |= (1 << 17) // Set self-contained bit
	}

	header[0] = byte(headerInt)
	header[1] = byte(headerInt >> 8)
	header[2] = byte(headerInt >> 16)

	// Calculate CRC24 for first 3 bytes
	crc24 := Crc24(header[:3])
	header[3] = byte(crc24)
	header[4] = byte(crc24 >> 8)
	header[5] = byte(crc24 >> 16)

	// Write header
	if _, err := w.Write(header[:]); err != nil {
		return fmt.Errorf("failed to write segment header: %w", err)
	}

	// Write payload
	if _, err := w.Write(s.Payload); err != nil {
		return fmt.Errorf("failed to write segment payload: %w", err)
	}

	// Calculate and write CRC32
	crc32 := Crc32(s.Payload)
	crc32Bytes := make([]byte, 4)
	binary.LittleEndian.PutUint32(crc32Bytes, crc32)
	if _, err := w.Write(crc32Bytes); err != nil {
		return fmt.Errorf("failed to write CRC32: %w", err)
	}

	return nil
}

func (s *Segment) writeCompressedSegment(w io.Writer, compressor gocql.Compressor) error {
	if compressor == nil {
		return fmt.Errorf("compressor required for compressed segment")
	}

	uncompressedLen := len(s.Payload)
	if uncompressedLen > MaxSegmentPayload {
		return fmt.Errorf("payload too large: %d bytes (max %d)", uncompressedLen, MaxSegmentPayload)
	}

	// Compress payload
	compressedPayload, err := compressor.AppendCompressed(nil, s.Payload)
	if err != nil {
		return fmt.Errorf("failed to compress payload: %w", err)
	}

	compressedLen := len(compressedPayload)

	// LZ4 returns 0 when the input is not compressible; send uncompressed.
	if compressedLen == 0 {
		compressedPayload = s.Payload
		compressedLen = uncompressedLen
		uncompressedLen = 0
	}

	// If compression doesn't reduce size, send uncompressed with uncompressed_len = 0
	if compressedLen >= uncompressedLen {
		compressedPayload = s.Payload
		compressedLen = uncompressedLen
		uncompressedLen = 0
	}

	// Build header (5 bytes)
	var headerBuf [8]byte
	combined := uint64(compressedLen) | uint64(uncompressedLen)<<17
	if s.SelfContained {
		combined |= (1 << 34) // Set self-contained bit
	}

	headerBuf[0] = byte(combined)
	headerBuf[1] = byte(combined >> 8)
	headerBuf[2] = byte(combined >> 16)
	headerBuf[3] = byte(combined >> 24)
	headerBuf[4] = byte(combined >> 32)

	// Calculate CRC24 for first 5 bytes
	crc24 := Crc24(headerBuf[:5])
	headerBuf[5] = byte(crc24)
	headerBuf[6] = byte(crc24 >> 8)
	headerBuf[7] = byte(crc24 >> 16)

	// Write header (first 5 bytes, then CRC24)
	if _, err := w.Write(headerBuf[:5]); err != nil {
		return fmt.Errorf("failed to write segment header: %w", err)
	}
	if _, err := w.Write(headerBuf[5:8]); err != nil {
		return fmt.Errorf("failed to write segment CRC24: %w", err)
	}

	// Write compressed payload
	if _, err := w.Write(compressedPayload); err != nil {
		return fmt.Errorf("failed to write segment payload: %w", err)
	}

	// Calculate and write CRC32
	crc32 := Crc32(compressedPayload)
	crc32Bytes := make([]byte, 4)
	binary.LittleEndian.PutUint32(crc32Bytes, crc32)
	if _, err := w.Write(crc32Bytes); err != nil {
		return fmt.Errorf("failed to write CRC32: %w", err)
	}

	return nil
}

// SegmentReader reads frames from segments (handles protocol v5+)
type SegmentReader struct {
	reader       io.Reader
	compressor   gocql.Compressor
	buffer       bytes.Buffer
	pendingFrame []byte // Partial frame from previous segment
}

// NewSegmentReader creates a new segment reader
func NewSegmentReader(r io.Reader, compressor gocql.Compressor) *SegmentReader {
	return &SegmentReader{
		reader:     r,
		compressor: compressor,
	}
}

// ReadFrame reads a complete frame from segments
func (sr *SegmentReader) ReadFrame() ([]byte, error) {
	// If we have a pending frame, return it
	if len(sr.pendingFrame) > 0 {
		frame := sr.pendingFrame
		sr.pendingFrame = nil
		return frame, nil
	}
	if sr.buffer.Len() >= FrameHeaderSize {
		if frame, err := sr.extractFrame(); err == nil {
			return frame, nil
		}
	}

	// Read segments until we get a complete frame or self-contained segment
	for {
		segment, err := ReadSegment(sr.reader, sr.compressor)
		if err != nil {
			return nil, err
		}

		// Add segment payload to buffer
		sr.buffer.Write(segment.Payload)

		if segment.SelfContained {
			// Parse frames from buffer
			return sr.extractFrame()
		}

		// Not self-contained, keep reading segments
		// We'll extract frames when we get a self-contained segment or have enough data
		if sr.buffer.Len() >= FrameHeaderSize {
			// Try to extract a frame
			frame, err := sr.extractFrame()
			if err == nil {
				return frame, nil
			}
			// If we can't extract a complete frame, continue reading
		}
	}
}

func (sr *SegmentReader) extractFrame() ([]byte, error) {
	if sr.buffer.Len() < FrameHeaderSize {
		return nil, fmt.Errorf("not enough data for frame header")
	}

	// Read header to get length
	header := sr.buffer.Bytes()[:FrameHeaderSize]
	length := int32(binary.BigEndian.Uint32(header[5:9]))

	// Check if we have complete frame
	totalSize := FrameHeaderSize + int(length)
	if sr.buffer.Len() < totalSize {
		return nil, fmt.Errorf("incomplete frame: need %d bytes, have %d", totalSize, sr.buffer.Len())
	}

	// Extract frame
	frame := make([]byte, totalSize)
	copy(frame, sr.buffer.Bytes()[:totalSize])

	// Remove from buffer
	remaining := sr.buffer.Bytes()[totalSize:]
	sr.buffer.Reset()
	sr.buffer.Write(remaining)

	return frame, nil
}

// SegmentWriter writes frames as segments (handles protocol v5+)
type SegmentWriter struct {
	writer     io.Writer
	compressor gocql.Compressor
}

// NewSegmentWriter creates a new segment writer
func NewSegmentWriter(w io.Writer, compressor gocql.Compressor) *SegmentWriter {
	return &SegmentWriter{
		writer:     w,
		compressor: compressor,
	}
}

// WriteFrame writes a frame as one or more segments
func (sw *SegmentWriter) WriteFrame(frame []byte, compress bool) error {
	if len(frame) == 0 {
		return nil
	}

	// If frame fits in one segment, write as self-contained
	if len(frame) <= MaxSegmentPayload {
		segment := &Segment{
			Payload:       frame,
			SelfContained: true,
			Compressed:    compress,
		}
		return segment.WriteSegment(sw.writer, sw.compressor)
	}

	// Split into multiple segments.
	// For a fragmented frame, all slices must be non-self-contained; the receiver
	// reconstructs completion from the frame header/body length in the first slice.
	offset := 0
	for offset < len(frame) {
		chunkSize := len(frame) - offset
		if chunkSize > MaxSegmentPayload {
			chunkSize = MaxSegmentPayload
		}

		segment := &Segment{
			Payload:       frame[offset : offset+chunkSize],
			SelfContained: false,
			Compressed:    compress,
		}

		if err := segment.WriteSegment(sw.writer, sw.compressor); err != nil {
			return err
		}

		offset += chunkSize
	}

	return nil
}