//go:build integration
// +build integration

package integration_test

import (
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

type loadWorkloadKey struct {
	userID      int
	sessionID   int
	sessionType int
}

type loadWorkloadRow struct {
	unreadCount   int64
	mentionCount  int64
	lastReadID    int64
	lastMessageID int64
	seqNo         int64
	archived      bool
}

func TestLoadStyleMixedWorkloadDataCorrectness(t *testing.T) {
	session := newSession(t, sessionOptions{
		disableHostLookup: true,
		numConns:          4,
		queryTimeout:      10 * time.Second,
	})

	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 := "load_style_chat_list"
	createTable := fmt.Sprintf(`
CREATE TABLE %s.%s (
	user_id int,
	session_id int,
	session_type int,
	scenario int,
	is_archived boolean,
	is_pinned boolean,
	is_muted boolean,
	unread_count int,
	mention_count int,
	last_read_id bigint,
	last_message_id bigint,
	last_message_time timestamp,
	content text,
	seq_no bigint,
	PRIMARY KEY (user_id, session_id, session_type)
)`, keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	const (
		scenario            = 1
		sessionType         = 1
		insertUserBase      = 1000
		selectUserBase      = 2000
		updateUserBase      = 3000
		usersPerRange       = 4
		seedSessionsPerUser = 4
		workers             = 4
		operationsPerWorker = 36
	)

	insertSQL := fmt.Sprintf(`
INSERT INTO %s.%s (
	user_id, session_id, session_type, scenario, is_archived, is_pinned, is_muted,
	unread_count, mention_count, last_read_id, last_message_id, last_message_time, content, seq_no
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, keyspace, table)

	type shadowStore struct {
		mu   sync.Mutex
		rows map[loadWorkloadKey]loadWorkloadRow
	}
	shadow := shadowStore{rows: map[loadWorkloadKey]loadWorkloadRow{}}
	putShadow := func(key loadWorkloadKey, row loadWorkloadRow) {
		shadow.mu.Lock()
		defer shadow.mu.Unlock()
		shadow.rows[key] = row
	}
	addShadow := func(key loadWorkloadKey, unreadDelta, mentionDelta, seqDelta, messageID int64) {
		shadow.mu.Lock()
		defer shadow.mu.Unlock()
		row := shadow.rows[key]
		row.unreadCount += unreadDelta
		row.mentionCount += mentionDelta
		row.seqNo += seqDelta
		row.lastMessageID = messageID
		shadow.rows[key] = row
	}
	markReadShadow := func(key loadWorkloadKey, lastReadID int64) {
		shadow.mu.Lock()
		defer shadow.mu.Unlock()
		row := shadow.rows[key]
		row.lastReadID = lastReadID
		row.unreadCount = 0
		row.mentionCount = 0
		row.archived = false
		shadow.rows[key] = row
	}
	archiveShadow := func(key loadWorkloadKey) {
		shadow.mu.Lock()
		defer shadow.mu.Unlock()
		row := shadow.rows[key]
		row.archived = true
		shadow.rows[key] = row
	}
	getUserTotals := func(userID int) (count, unread, mention int64) {
		shadow.mu.Lock()
		defer shadow.mu.Unlock()
		for key, row := range shadow.rows {
			if key.userID != userID || key.sessionType != sessionType || row.archived {
				continue
			}
			count++
			unread += row.unreadCount
			mention += row.mentionCount
		}
		return count, unread, mention
	}

	for _, base := range []int{selectUserBase, updateUserBase} {
		for userOffset := 0; userOffset < usersPerRange; userOffset++ {
			userID := base + userOffset
			for sessionOffset := 0; sessionOffset < seedSessionsPerUser; sessionOffset++ {
				key := loadWorkloadKey{userID: userID, sessionID: 10_000 + userOffset*100 + sessionOffset, sessionType: sessionType}
				row := loadWorkloadRow{
					unreadCount:   int64(sessionOffset + 1),
					mentionCount:  int64(sessionOffset % 2),
					lastMessageID: int64(1_000_000 + userID*100 + sessionOffset),
					seqNo:         int64(sessionOffset),
				}
				if err := session.Query(
					insertSQL,
					key.userID,
					key.sessionID,
					key.sessionType,
					scenario,
					false,
					sessionOffset == 0,
					false,
					int(row.unreadCount),
					int(row.mentionCount),
					row.lastReadID,
					row.lastMessageID,
					time.Unix(int64(sessionOffset), 0).UTC(),
					fmt.Sprintf("seed-%d-%d", userID, sessionOffset),
					row.seqNo,
				).Exec(); err != nil {
					t.Fatalf("seed row %+v failed: %v", key, err)
				}
				putShadow(key, row)
			}
		}
	}

	updateSQL := fmt.Sprintf(`
UPDATE %s.%s
SET unread_count = unread_count + ?,
    mention_count = mention_count + ?,
    seq_no = seq_no + ?,
    last_message_id = ?
WHERE user_id = ? AND session_id = ? AND session_type = ?`, keyspace, table)
	var nextInsertSession atomic.Int64
	nextInsertSession.Store(50_000)
	var nextMessageID atomic.Int64
	nextMessageID.Store(9_000_000)

	var wg sync.WaitGroup
	errCh := make(chan error, workers)
	for worker := 0; worker < workers; worker++ {
		wg.Add(1)
		go func(worker int) {
			defer wg.Done()
			for op := 0; op < operationsPerWorker; op++ {
				switch op % 10 {
				case 0, 1, 2, 3, 4, 5:
					batch := session.NewBatch(gocql.UnloggedBatch)
					keys := make([]loadWorkloadKey, 0, 2)
					rows := make([]loadWorkloadRow, 0, 2)
					for i := 0; i < 2; i++ {
						userID := insertUserBase + (worker+i)%usersPerRange
						sessionID := int(nextInsertSession.Add(1))
						messageID := nextMessageID.Add(1)
						key := loadWorkloadKey{userID: userID, sessionID: sessionID, sessionType: sessionType}
						row := loadWorkloadRow{
							unreadCount:   1,
							mentionCount:  int64((worker + op + i) % 2),
							lastMessageID: messageID,
							seqNo:         messageID,
						}
						batch.Query(
							insertSQL,
							key.userID,
							key.sessionID,
							key.sessionType,
							scenario,
							false,
							false,
							false,
							int(row.unreadCount),
							int(row.mentionCount),
							row.lastReadID,
							row.lastMessageID,
							time.Unix(messageID, 0).UTC(),
							fmt.Sprintf("fanout-%d-%d", worker, op),
							row.seqNo,
						)
						keys = append(keys, key)
						rows = append(rows, row)
					}
					if err := session.ExecuteBatch(batch); err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d insert batch failed: %w", worker, op, err)
						return
					}
					for i, key := range keys {
						putShadow(key, rows[i])
					}
				case 6, 7:
					userID := selectUserBase + (worker+op)%usersPerRange
					wantCount, wantUnread, wantMention := getUserTotals(userID)
					row, err := singleRow(session, fmt.Sprintf(
						"SELECT COUNT(*) AS cnt, SUM(unread_count) AS sum_unread, SUM(mention_count) AS sum_mention FROM %s.%s WHERE user_id = %d AND scenario = %d AND is_archived = 0 ALLOW FILTERING",
						keyspace,
						table,
						userID,
						scenario,
					))
					if err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d aggregate select failed: %w", worker, op, err)
						return
					}
					gotCount, err := parseIntField(row["cnt"])
					if err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d parse cnt failed: %w", worker, op, err)
						return
					}
					gotUnread, err := parseNullableIntField(row["sum_unread"])
					if err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d parse sum_unread failed: %w", worker, op, err)
						return
					}
					gotMention, err := parseNullableIntField(row["sum_mention"])
					if err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d parse sum_mention failed: %w", worker, op, err)
						return
					}
					if gotCount != wantCount || gotUnread != wantUnread || gotMention != wantMention {
						errCh <- fmt.Errorf(
							"worker=%d op=%d select shadow mismatch user=%d got count=%d unread=%d mention=%d want count=%d unread=%d mention=%d",
							worker,
							op,
							userID,
							gotCount,
							gotUnread,
							gotMention,
							wantCount,
							wantUnread,
							wantMention,
						)
						return
					}
					if _, err := scanRows(session, fmt.Sprintf(
						"SELECT session_id, unread_count, mention_count FROM %s.%s WHERE user_id = %d AND scenario = %d AND is_archived = 0 ORDER BY last_message_time DESC LIMIT 20 ALLOW FILTERING",
						keyspace,
						table,
						userID,
						scenario,
					)); err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d list select failed: %w", worker, op, err)
						return
					}
				default:
					userID := updateUserBase + (worker+op)%usersPerRange
					sessionID := 10_000 + ((worker + op) % usersPerRange * 100) + (op % seedSessionsPerUser)
					key := loadWorkloadKey{userID: userID, sessionID: sessionID, sessionType: sessionType}
					unreadDelta := int64(1 + (worker+op)%3)
					mentionDelta := int64((worker + op) % 2)
					messageID := nextMessageID.Add(1)
					if err := session.Query(updateSQL, int(unreadDelta), int(mentionDelta), int64(1), messageID, key.userID, key.sessionID, key.sessionType).Exec(); err != nil {
						errCh <- fmt.Errorf("worker=%d op=%d update failed: %w", worker, op, err)
						return
					}
					addShadow(key, unreadDelta, mentionDelta, 1, messageID)
				}
			}
		}(worker)
	}

	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatal(err)
		}
	}

	verifyLoadWorkloadShadow(t, session, keyspace, table, shadow.rows)

	markReadKey := loadWorkloadKey{userID: updateUserBase, sessionID: 10_000, sessionType: sessionType}
	lastReadID := nextMessageID.Add(1)
	markReadSQL := fmt.Sprintf(`
UPDATE %s.%s
SET last_read_id = ?, unread_count = 0, mention_count = 0, is_archived = 0
WHERE user_id = ? AND session_id = ? AND session_type = ?`, keyspace, table)
	if err := session.Query(markReadSQL, lastReadID, markReadKey.userID, markReadKey.sessionID, markReadKey.sessionType).Exec(); err != nil {
		t.Fatalf("mark read update failed: %v", err)
	}
	markReadShadow(markReadKey, lastReadID)

	archiveKey := loadWorkloadKey{userID: updateUserBase + 1, sessionID: 10_100, sessionType: sessionType}
	archiveSQL := fmt.Sprintf("UPDATE %s.%s SET is_archived = 1 WHERE user_id = ? AND session_id = ? AND session_type = ?", keyspace, table)
	if err := session.Query(archiveSQL, archiveKey.userID, archiveKey.sessionID, archiveKey.sessionType).Exec(); err != nil {
		t.Fatalf("archive update failed: %v", err)
	}
	archiveShadow(archiveKey)

	verifyLoadWorkloadShadow(t, session, keyspace, table, shadow.rows)

	for _, userID := range []int{selectUserBase, updateUserBase, insertUserBase} {
		row, err := singleRow(session, fmt.Sprintf(
			"SELECT COUNT(*) AS cnt, SUM(unread_count) AS sum_unread, SUM(mention_count) AS sum_mention FROM %s.%s WHERE user_id = %d AND scenario = %d AND is_archived = 0 ALLOW FILTERING",
			keyspace,
			table,
			userID,
			scenario,
		))
		if err != nil {
			t.Fatalf("final aggregate for user %d failed: %v", userID, err)
		}
		count, unread, mention := expectedLoadWorkloadTotals(shadow.rows, userID, sessionType)
		gotCount, err := parseIntField(row["cnt"])
		if err != nil {
			t.Fatalf("parse final count for user %d failed: %v", userID, err)
		}
		gotUnread, err := parseNullableIntField(row["sum_unread"])
		if err != nil {
			t.Fatalf("parse final unread sum for user %d failed: %v", userID, err)
		}
		gotMention, err := parseNullableIntField(row["sum_mention"])
		if err != nil {
			t.Fatalf("parse final mention sum for user %d failed: %v", userID, err)
		}
		if gotCount != count || gotUnread != unread || gotMention != mention {
			t.Fatalf("final aggregate mismatch user=%d got count=%d unread=%d mention=%d want count=%d unread=%d mention=%d", userID, gotCount, gotUnread, gotMention, count, unread, mention)
		}
	}

	for _, userID := range []int{selectUserBase, updateUserBase, insertUserBase} {
		if _, err := scanRows(session, fmt.Sprintf(
			"SELECT session_id, unread_count, mention_count FROM %s.%s WHERE user_id = %d AND scenario = %d AND is_archived = 0 ORDER BY last_message_time DESC LIMIT 20 ALLOW FILTERING",
			keyspace,
			table,
			userID,
			scenario,
		)); err != nil {
			t.Fatalf("final list select for user %d failed: %v", userID, err)
		}
	}
}

func verifyLoadWorkloadShadow(t *testing.T, session *gocql.Session, keyspace, table string, expected map[loadWorkloadKey]loadWorkloadRow) {
	t.Helper()
	for key, want := range expected {
		row, err := singleRow(session, fmt.Sprintf(
			"SELECT unread_count, mention_count, last_read_id, last_message_id, seq_no, is_archived FROM %s.%s WHERE user_id = %d AND session_id = %d AND session_type = %d",
			keyspace,
			table,
			key.userID,
			key.sessionID,
			key.sessionType,
		))
		if err != nil {
			t.Fatalf("select row %+v failed: %v", key, err)
		}
		unread, err := parseIntField(row["unread_count"])
		if err != nil {
			t.Fatalf("parse unread_count for %+v failed: %v", key, err)
		}
		mention, err := parseIntField(row["mention_count"])
		if err != nil {
			t.Fatalf("parse mention_count for %+v failed: %v", key, err)
		}
		lastRead, err := parseNullableIntField(row["last_read_id"])
		if err != nil {
			t.Fatalf("parse last_read_id for %+v failed: %v", key, err)
		}
		lastMessage, err := parseIntField(row["last_message_id"])
		if err != nil {
			t.Fatalf("parse last_message_id for %+v failed: %v", key, err)
		}
		seqNo, err := parseIntField(row["seq_no"])
		if err != nil {
			t.Fatalf("parse seq_no for %+v failed: %v", key, err)
		}
		archived, ok := row["is_archived"].(bool)
		if !ok {
			t.Fatalf("parse is_archived for %+v failed: got %T %#v", key, row["is_archived"], row["is_archived"])
		}
		if unread != want.unreadCount || mention != want.mentionCount || lastRead != want.lastReadID || lastMessage != want.lastMessageID || seqNo != want.seqNo || archived != want.archived {
			t.Fatalf(
				"load workload shadow mismatch for %+v: got unread=%d mention=%d lastRead=%d lastMessage=%d seq=%d archived=%v, want unread=%d mention=%d lastRead=%d lastMessage=%d seq=%d archived=%v",
				key,
				unread,
				mention,
				lastRead,
				lastMessage,
				seqNo,
				archived,
				want.unreadCount,
				want.mentionCount,
				want.lastReadID,
				want.lastMessageID,
				want.seqNo,
				want.archived,
			)
		}
	}
}

func expectedLoadWorkloadTotals(rows map[loadWorkloadKey]loadWorkloadRow, userID, sessionType int) (count, unread, mention int64) {
	for key, row := range rows {
		if key.userID != userID || key.sessionType != sessionType || row.archived {
			continue
		}
		count++
		unread += row.unreadCount
		mention += row.mentionCount
	}
	return count, unread, mention
}

func parseNullableIntField(value interface{}) (int64, error) {
	if value == nil {
		return 0, nil
	}
	return parseIntField(value)
}