//go:build integration
// +build integration

package integration_test

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/executor"
)

type rawProtocolOptions struct {
	username string
	password string
}

type rawProtocolConn struct {
	t          *testing.T
	conn       net.Conn
	version    byte
	useSegment bool
	reader     *bufio.Reader
	segReader  *cql.FrameReader
	segWriter  *cql.FrameWriter
}

func newRawProtocolConn(t *testing.T, version byte, opts rawProtocolOptions) *rawProtocolConn {
	t.Helper()
	requireProxy(t)

	addr := fmt.Sprintf("%s:%d", proxyHost, proxyPort)
	conn, err := net.DialTimeout("tcp", addr, 5*time.Second)
	if err != nil {
		t.Fatalf("dial proxy failed: %v", err)
	}

	client := &rawProtocolConn{
		t:       t,
		conn:    conn,
		version: version,
		reader:  bufio.NewReader(conn),
	}
	t.Cleanup(func() {
		_ = conn.Close()
	})

	client.startup(opts)
	return client
}

func (c *rawProtocolConn) startup(opts rawProtocolOptions) {
	startup := cql.BuildStartupFrame("3.4.0", "fidelity-test", "1.0.0", "")
	startup.Version = c.version
	c.writeClassic(startup)

	resp := c.readClassic()
	switch resp.Opcode {
	case cql.OpcodeAuthenticate:
		if c.version >= cql.ProtoVersion5 {
			c.enableSegments()
		}
		username := opts.username
		if username == "" {
			username = mysqlCfg.user
		}
		password := opts.password
		if password == "" {
			password = mysqlCfg.password
		}
		auth := cql.BuildAuthResponseFrame(username, password)
		auth.Version = c.version
		c.writeFrame(auth)
		authResp := c.readFrame()
		if authResp.Opcode != cql.OpcodeAuthSuccess {
			c.t.Fatalf("expected AUTH_SUCCESS, got opcode %d", authResp.Opcode)
		}
	case cql.OpcodeReady:
		if c.version >= cql.ProtoVersion5 {
			c.enableSegments()
		}
	default:
		c.t.Fatalf("unexpected startup response opcode %d", resp.Opcode)
	}
}

func (c *rawProtocolConn) enableSegments() {
	if c.useSegment {
		return
	}
	c.useSegment = true
	c.segReader = cql.NewFrameReader(c.reader, nil, c.version)
	c.segWriter = cql.NewFrameWriter(c.conn, nil, c.version)
}

func (c *rawProtocolConn) writeClassic(frame *cql.Frame) {
	c.t.Helper()
	frame.Version = c.version
	_ = c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	if err := frame.WriteFrame(c.conn); err != nil {
		c.t.Fatalf("write frame failed: %v", err)
	}
}

func (c *rawProtocolConn) writeFrame(frame *cql.Frame) {
	c.t.Helper()
	frame.Version = c.version
	_ = c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	if c.useSegment {
		if err := c.segWriter.WriteFrame(frame); err != nil {
			c.t.Fatalf("write segment frame failed: %v", err)
		}
		return
	}
	if err := frame.WriteFrame(c.conn); err != nil {
		c.t.Fatalf("write frame failed: %v", err)
	}
}

func (c *rawProtocolConn) readClassic() *cql.Frame {
	c.t.Helper()
	_ = c.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	frame, err := cql.ReadFrame(c.reader)
	if err != nil {
		c.t.Fatalf("read frame failed: %v", err)
	}
	return frame
}

func (c *rawProtocolConn) readFrame() *cql.Frame {
	c.t.Helper()
	_ = c.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	if c.useSegment {
		frame, err := c.segReader.ReadFrame()
		if err != nil {
			c.t.Fatalf("read segment frame failed: %v", err)
		}
		return frame
	}
	frame, err := cql.ReadFrame(c.reader)
	if err != nil {
		c.t.Fatalf("read frame failed: %v", err)
	}
	return frame
}

func TestProtocolFrameHandshakeV3ToV5(t *testing.T) {
	versions := []byte{cql.ProtoVersion3, cql.ProtoVersion4, cql.ProtoVersion5}
	for _, version := range versions {
		t.Run(fmt.Sprintf("v%d", version), func(t *testing.T) {
			client := newRawProtocolConn(t, version, rawProtocolOptions{})
			options := cql.BuildOptionsFrame()
			client.writeFrame(options)
			resp := client.readFrame()
			if resp.Opcode != cql.OpcodeSupported {
				t.Fatalf("expected SUPPORTED, got opcode %d", resp.Opcode)
			}
			opts, err := parseSupportedOptions(resp.Body)
			if err != nil {
				t.Fatalf("parse supported options failed: %v", err)
			}
			if _, ok := opts["CQL_VERSION"]; !ok {
				t.Fatalf("expected CQL_VERSION in supported options")
			}
		})
	}
}

func TestProtocolInvalidQueryFlags(t *testing.T) {
	client := newRawProtocolConn(t, cql.ProtoVersion4, rawProtocolOptions{})

	cases := []struct {
		name  string
		flags uint32
		vals  []cql.QueryValue
	}{
		{
			name:  "named values",
			flags: cql.QueryFlagValues | cql.QueryFlagWithNameValues,
			vals:  []cql.QueryValue{{Name: "id", Value: encodeInt32(1)}},
		},
		{
			name:  "positional values",
			flags: cql.QueryFlagValues,
			vals:  []cql.QueryValue{{Value: encodeInt32(1)}},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			frame := cql.BuildQueryFrame("SELECT key FROM system.local", cql.ConsistencyOne, tc.flags, tc.vals, 0, nil, 0, 0, "", nil)
			frame.Stream = 1
			client.writeFrame(frame)
			resp := client.readFrame()
			code := parseErrorCode(t, resp)
			if code != cql.ErrorCodeInvalid {
				t.Fatalf("expected ErrorCodeInvalid, got %d", code)
			}
		})
	}
}

func TestProtocolExecuteNamedValuesInvalid(t *testing.T) {
	client := newRawProtocolConn(t, cql.ProtoVersion4, rawProtocolOptions{})

	prepare := buildPrepareFrame(client.version, "SELECT key FROM system.local WHERE key = ?", "")
	prepare.Stream = 1
	client.writeFrame(prepare)
	prepared := client.readFrame()
	stmtID := parsePreparedID(t, prepared)

	values := []cql.QueryValue{{Name: "key", Value: encodeText("local")}}
	execute := cql.BuildExecuteFrame(stmtID, nil, cql.ConsistencyOne, cql.QueryFlagValues|cql.QueryFlagWithNameValues, values, 0, nil, 0, 0, "", nil)
	execute.Stream = 2
	client.writeFrame(execute)
	resp := client.readFrame()
	code := parseErrorCode(t, resp)
	if code != cql.ErrorCodeInvalid {
		t.Fatalf("expected ErrorCodeInvalid, got %d", code)
	}
}

func TestProtocolStreamMultiplexingV4(t *testing.T) {
	client := newRawProtocolConn(t, cql.ProtoVersion4, rawProtocolOptions{})

	const requests = 50
	expected := make(map[int16]struct{}, requests)
	for i := 0; i < requests; i++ {
		stream := int16(i + 1)
		frame := cql.BuildQueryFrame("SELECT key FROM system.local", cql.ConsistencyOne, 0, nil, 0, nil, 0, 0, "", nil)
		frame.Stream = stream
		client.writeFrame(frame)
		expected[stream] = struct{}{}
	}

	for i := 0; i < requests; i++ {
		resp := client.readFrame()
		if resp.Opcode != cql.OpcodeResult {
			t.Fatalf("expected RESULT, got opcode %d", resp.Opcode)
		}
		if _, ok := expected[resp.Stream]; !ok {
			t.Fatalf("unexpected stream id %d", resp.Stream)
		}
		delete(expected, resp.Stream)
	}
	if len(expected) != 0 {
		t.Fatalf("missing responses for %d streams", len(expected))
	}
}

func TestProtocolPreparedReuseSoak(t *testing.T) {
	client := newRawProtocolConn(t, cql.ProtoVersion4, rawProtocolOptions{})

	prepare := buildPrepareFrame(client.version, "SELECT key FROM system.local WHERE key = ?", "")
	prepare.Stream = 1
	client.writeFrame(prepare)
	prepared := client.readFrame()
	stmtID := parsePreparedID(t, prepared)

	values := []cql.QueryValue{{Value: encodeText("local")}}
	for i := 0; i < 25; i++ {
		execute := cql.BuildExecuteFrame(stmtID, nil, cql.ConsistencyOne, cql.QueryFlagValues, values, 0, nil, 0, 0, "", nil)
		execute.Stream = int16(i + 2)
		client.writeFrame(execute)
		resp := client.readFrame()
		if resp.Opcode != cql.OpcodeResult {
			t.Fatalf("expected RESULT, got opcode %d", resp.Opcode)
		}
	}
}

func TestProtocolPagingStateIgnored(t *testing.T) {
	session := newSession(t, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	table := "paging_items"
	if err := session.Query(createTableCQL(keyspace, table)).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}
	insert := fmt.Sprintf("INSERT INTO %s.%s (id, name, idx) VALUES (?, ?, ?)", keyspace, table)
	for i := 0; i < 3; i++ {
		if err := session.Query(insert, i, fmt.Sprintf("name-%d", i), i).Exec(); err != nil {
			t.Fatalf("insert failed: %v", err)
		}
	}

	client := newRawProtocolConn(t, cql.ProtoVersion4, rawProtocolOptions{})
	flags := uint32(cql.QueryFlagPageSize | cql.QueryFlagWithPagingState)
	frame := cql.BuildQueryFrame(fmt.Sprintf("SELECT id FROM %s.%s", keyspace, table), cql.ConsistencyOne, flags, nil, 1, []byte("state"), 0, 0, "", nil)
	frame.Stream = 5
	client.writeFrame(frame)
	resp := client.readFrame()
	if resp.Opcode != cql.OpcodeResult {
		t.Fatalf("expected RESULT, got opcode %d", resp.Opcode)
	}
	flagsVal := parseRowsMetadataFlags(t, resp)
	if flagsVal&resultMetadataHasMorePages != 0 {
		t.Fatalf("expected no paging state flag, got 0x%x", flagsVal)
	}
}

func TestProtocolKeyspaceFlagsV5PrepareExecute(t *testing.T) {
	session := newSession(t, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	table := "ks_flags"
	if err := session.Query(createTableCQL(keyspace, table)).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	client := newRawProtocolConn(t, cql.ProtoVersion5, rawProtocolOptions{})

	t.Run("prepare with keyspace", func(t *testing.T) {
		prepare := buildPrepareFrame(client.version, fmt.Sprintf("INSERT INTO %s (id, name, idx) VALUES (?, ?, ?)", table), keyspace)
		prepare.Stream = 1
		client.writeFrame(prepare)
		prepared := client.readFrame()
		stmtID := parsePreparedID(t, prepared)

		execValues := [][]byte{encodeInt32(1), encodeText("alpha"), encodeInt32(1)}
		execute := buildExecuteFrame(client.version, stmtID, nil, cql.ConsistencyOne, cql.QueryFlagValues, execValues, 0, "")
		execute.Stream = 2
		client.writeFrame(execute)
		resp := client.readFrame()
		if parseResultKind(t, resp) != executor.ResultKindVoid {
			t.Fatalf("expected VOID result")
		}

		var got string
		query := fmt.Sprintf("SELECT name FROM %s.%s WHERE id = ?", keyspace, table)
		if err := session.Query(query, 1).Scan(&got); err != nil {
			t.Fatalf("select failed: %v", err)
		}
		if got != "alpha" {
			t.Fatalf("unexpected name: %q", got)
		}
	})

	t.Run("execute with keyspace", func(t *testing.T) {
		prepare := buildPrepareFrame(client.version, fmt.Sprintf("INSERT INTO %s (id, name, idx) VALUES (?, ?, ?)", table), "")
		prepare.Stream = 3
		client.writeFrame(prepare)
		prepared := client.readFrame()
		stmtID := parsePreparedID(t, prepared)

		execValues := [][]byte{encodeInt32(2), encodeText("beta"), encodeInt32(2)}
		flags := uint32(cql.QueryFlagValues | cql.QueryFlagWithKeyspace)
		execute := buildExecuteFrame(client.version, stmtID, nil, cql.ConsistencyOne, flags, execValues, 0, keyspace)
		execute.Stream = 4
		client.writeFrame(execute)
		resp := client.readFrame()
		if parseResultKind(t, resp) != executor.ResultKindVoid {
			t.Fatalf("expected VOID result")
		}

		var got string
		query := fmt.Sprintf("SELECT name FROM %s.%s WHERE id = ?", keyspace, table)
		if err := session.Query(query, 2).Scan(&got); err != nil {
			t.Fatalf("select failed: %v", err)
		}
		if got != "beta" {
			t.Fatalf("unexpected name: %q", got)
		}
	})
}

const resultMetadataHasMorePages = 0x0002

func parseRowsMetadataFlags(t *testing.T, frame *cql.Frame) int32 {
	t.Helper()
	if frame.Opcode != cql.OpcodeResult {
		t.Fatalf("expected RESULT frame, got %d", frame.Opcode)
	}
	if len(frame.Body) < 8 {
		t.Fatalf("result body too short")
	}
	parser := cql.NewFrameBodyParser(frame.Body)
	kind, err := parser.ReadInt()
	if err != nil {
		t.Fatalf("read result kind failed: %v", err)
	}
	if kind != int32(cql.ResultKindRows) {
		t.Fatalf("expected ROWS result, got %d", kind)
	}
	flags, err := parser.ReadInt()
	if err != nil {
		t.Fatalf("read metadata flags failed: %v", err)
	}
	return flags
}

func parseResultKind(t *testing.T, frame *cql.Frame) executor.ResultKind {
	t.Helper()
	if frame.Opcode != cql.OpcodeResult {
		t.Fatalf("expected RESULT frame, got %d", frame.Opcode)
	}
	if len(frame.Body) < 4 {
		t.Fatalf("result body too short")
	}
	kind := executor.ResultKind(binary.BigEndian.Uint32(frame.Body[:4]))
	return kind
}

func parsePreparedID(t *testing.T, frame *cql.Frame) []byte {
	t.Helper()
	if frame.Opcode == cql.OpcodeError {
		code, msg := parseError(t, frame)
		t.Fatalf("prepare failed: code=%d msg=%q", code, msg)
	}
	if frame.Opcode != cql.OpcodeResult {
		t.Fatalf("expected RESULT frame, got %d", frame.Opcode)
	}
	parser := cql.NewFrameBodyParser(frame.Body)
	kind, err := parser.ReadInt()
	if err != nil {
		t.Fatalf("read result kind failed: %v", err)
	}
	if kind != int32(cql.ResultKindPrepared) {
		t.Fatalf("expected PREPARED result, got %d", kind)
	}
	stmtID, err := parser.ReadShortBytes()
	if err != nil {
		t.Fatalf("read prepared id failed: %v", err)
	}
	if frame.Version >= cql.ProtoVersion5 {
		if _, err := parser.ReadShortBytes(); err != nil {
			t.Fatalf("read result metadata id failed: %v", err)
		}
	}
	return stmtID
}

func parseErrorCode(t *testing.T, frame *cql.Frame) int {
	t.Helper()
	code, _ := parseError(t, frame)
	return code
}

func parseError(t *testing.T, frame *cql.Frame) (int, string) {
	t.Helper()
	if frame.Opcode != cql.OpcodeError {
		t.Fatalf("expected ERROR frame, got %d", frame.Opcode)
	}
	parser := cql.NewFrameBodyParser(frame.Body)
	code, err := parser.ReadInt()
	if err != nil {
		t.Fatalf("read error code failed: %v", err)
	}
	message, err := parser.ReadString()
	if err != nil {
		t.Fatalf("read error message failed: %v", err)
	}
	return int(code), message
}

func parseSupportedOptions(body []byte) (map[string][]string, error) {
	parser := cql.NewFrameBodyParser(body)
	count, err := parser.ReadShort()
	if err != nil {
		return nil, err
	}
	result := make(map[string][]string, int(count))
	for i := 0; i < int(count); i++ {
		key, err := parser.ReadString()
		if err != nil {
			return nil, err
		}
		listCount, err := parser.ReadShort()
		if err != nil {
			return nil, err
		}
		values := make([]string, 0, listCount)
		for j := 0; j < int(listCount); j++ {
			val, err := parser.ReadString()
			if err != nil {
				return nil, err
			}
			values = append(values, val)
		}
		result[key] = values
	}
	return result, nil
}

func encodeInt32(v int32) []byte {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, uint32(v))
	return buf
}

func encodeText(v string) []byte {
	return []byte(v)
}

func buildExecuteFrame(version byte, preparedID []byte, resultMetadataID []byte, consistency uint16, flags uint32, values [][]byte, serialConsistency uint16, keyspace string) *cql.Frame {
	writer := cql.NewFrameBodyWriter()
	_ = writer.WriteShortBytes(preparedID)
	if version >= cql.ProtoVersion5 {
		if resultMetadataID == nil {
			resultMetadataID = []byte{}
		}
		_ = writer.WriteShortBytes(resultMetadataID)
	}
	_ = writer.WriteConsistency(consistency)
	_ = writer.WriteFlags(flags, version)
	if flags&cql.QueryFlagValues != 0 {
		_ = writer.WriteShort(uint16(len(values)))
		for _, val := range values {
			_ = writer.WriteBytes(val)
		}
	}
	if flags&cql.QueryFlagWithSerialConsistency != 0 {
		_ = writer.WriteConsistency(serialConsistency)
	}
	if flags&cql.QueryFlagWithKeyspace != 0 {
		_ = writer.WriteString(keyspace)
	}
	body := writer.Bytes()
	return &cql.Frame{Version: version, Opcode: cql.OpcodeExecute, Length: int32(len(body)), Body: body}
}

func buildQueryFrameWithVersion(version byte, query string, consistency uint16, flags uint32, values []cql.QueryValue, pageSize int, pagingState []byte, serialConsistency uint16, timestamp int64, keyspace string, nowInSeconds *int32) *cql.Frame {
	var buf bytes.Buffer

	buf.Write(writeLongString(query))
	buf.Write(writeConsistency(consistency))

	if version >= cql.ProtoVersion5 {
		buf.Write(writeInt(uint32(flags)))
	} else {
		buf.Write([]byte{byte(flags)})
	}

	if flags&cql.QueryFlagValues != 0 {
		buf.Write(writeShort(uint16(len(values))))
		for _, v := range values {
			if flags&cql.QueryFlagWithNameValues != 0 {
				buf.Write(writeString(v.Name))
			}
			writeQueryValue(&buf, v)
		}
	}

	if flags&cql.QueryFlagPageSize != 0 {
		buf.Write(writeInt(uint32(pageSize)))
	}
	if flags&cql.QueryFlagWithPagingState != 0 {
		buf.Write(writeBytes(pagingState))
	}
	if flags&cql.QueryFlagWithSerialConsistency != 0 {
		buf.Write(writeConsistency(serialConsistency))
	}
	if flags&cql.QueryFlagWithDefaultTimestamp != 0 {
		buf.Write(writeLong(timestamp))
	}
	if flags&cql.QueryFlagWithKeyspace != 0 {
		buf.Write(writeString(keyspace))
	}
	if flags&cql.QueryFlagNowInSeconds != 0 && nowInSeconds != nil {
		buf.Write(writeInt(uint32(*nowInSeconds)))
	}

	body := buf.Bytes()
	return &cql.Frame{Version: version, Opcode: cql.OpcodeQuery, Length: int32(len(body)), Body: body}
}

func buildPrepareFrame(version byte, statement string, keyspace string) *cql.Frame {
	var buf bytes.Buffer
	buf.Write(writeLongString(statement))

	if version >= cql.ProtoVersion5 {
		flags := uint32(0)
		if keyspace != "" {
			flags |= cql.PrepareFlagWithKeyspace
		}
		flagBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(flagBytes, flags)
		buf.Write(flagBytes)
		if keyspace != "" {
			buf.Write(writeString(keyspace))
		}
	}

	body := buf.Bytes()
	return &cql.Frame{Version: version, Opcode: cql.OpcodePrepare, Length: int32(len(body)), Body: body}
}

func writeLongString(value string) []byte {
	buf := make([]byte, 4+len(value))
	binary.BigEndian.PutUint32(buf[:4], uint32(len(value)))
	copy(buf[4:], value)
	return buf
}

func writeString(value string) []byte {
	buf := make([]byte, 2+len(value))
	binary.BigEndian.PutUint16(buf[:2], uint16(len(value)))
	copy(buf[2:], value)
	return buf
}

func writeShort(value uint16) []byte {
	buf := make([]byte, 2)
	binary.BigEndian.PutUint16(buf, value)
	return buf
}

func writeInt(value uint32) []byte {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, value)
	return buf
}

func writeLong(value int64) []byte {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, uint64(value))
	return buf
}

func writeConsistency(value uint16) []byte {
	return writeShort(value)
}

func writeBytes(value []byte) []byte {
	if value == nil {
		buf := make([]byte, 4)
		binary.BigEndian.PutUint32(buf, uint32(0xffffffff))
		return buf
	}
	buf := make([]byte, 4+len(value))
	binary.BigEndian.PutUint32(buf[:4], uint32(len(value)))
	copy(buf[4:], value)
	return buf
}

func writeQueryValue(buf *bytes.Buffer, value cql.QueryValue) {
	if value.Unset {
		buf.Write(writeInt(uint32(0xfffffffe)))
		return
	}
	buf.Write(writeBytes(value.Value))
}

func executeRawQuery(t *testing.T, client *rawProtocolConn, query string, stream int16) *cql.Frame {
	t.Helper()
	frame := buildQueryFrameWithVersion(client.version, query, cql.ConsistencyOne, 0, nil, 0, nil, 0, 0, "", nil)
	frame.Stream = stream
	client.writeFrame(frame)
	return client.readFrame()
}