package main

import (
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"igit.com/xbase/raft/db"
)

const (
	TotalKeys = 100000
	DataDir   = "bench_db_data"
)

func main() {
	// Scenario: Key Query Performance and Correctness
	fmt.Println("==================================================")
	fmt.Println("SCENARIO: Key Query Performance (Flat Array Index)")
	fmt.Println("==================================================")
	
	// Value Index Disabled
	runBenchmark(false)
}

func runBenchmark(enableValueIndex bool) {
	// Clean up previous run
	os.RemoveAll(DataDir)

	fmt.Printf("Initializing DB Engine in %s (ValueIndex=%v)...\n", DataDir, enableValueIndex)
	
	startMem := getMemUsage()
	
	e, err := db.NewEngine(DataDir, db.WithValueIndex(enableValueIndex))
	if err != nil {
		panic(err)
	}
	defer e.Close()
	defer os.RemoveAll(DataDir) // Cleanup after run

	// 1. Bulk Insert 100k keys with varied segments (1 to 5)
	fmt.Println("\n--- Phase 1: Bulk Insert 100k Keys (Varied Segments) ---")
	keys := make([]string, TotalKeys)
	start := time.Now()
	
	// Generate keys with 5 patterns
	// Pattern 1 (Len 1): item.<id>                      (20k)
	// Pattern 2 (Len 2): user.<id>.profile              (20k)
	// Pattern 3 (Len 3): team.<team_id>.user.<user_id> (20k) -> 200 users per team
	// Pattern 4 (Len 4): region.<r>.zone.<z>.node.<n>   (20k) -> 100 zones, 200 nodes
	// Pattern 5 (Len 5): a.b.c.d.<id>                   (20k)

	for i := 0; i < TotalKeys; i++ {
		group := i % 5
		id := i / 5
		
		switch group {
		case 0:
			// Len 1: item.<id>
			// But to test prefix well, let's group them slightly? 
			// Actually "item.1" is 2 segments if split by dot.
			// User said "one to 5 segments". 
			// Let's treat "." as separator.
			// "root<id>" is 1 segment.
			keys[i] = fmt.Sprintf("root%d", id)
		case 1:
			// Len 2: user.<id>
			keys[i] = fmt.Sprintf("user.%d", id)
		case 2:
			// Len 3: group.<id%100>.member.<id>
			// group.0.member.0 ... group.0.member.199
			keys[i] = fmt.Sprintf("group.%d.member.%d", id%100, id)
		case 3:
			// Len 4: app.<id%10>.ver.<id%10>.config.<id>
			keys[i] = fmt.Sprintf("app.%d.ver.%d.config.%d", id%10, id%10, id)
		case 4:
			// Len 5: log.2023.01.01.<id>
			keys[i] = fmt.Sprintf("log.2023.01.01.%d", id)
		}
	}

	var wg sync.WaitGroup
	workers := 10
	chunkSize := TotalKeys / workers
	var insertOps int64

	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			base := id * chunkSize
			for i := 0; i < chunkSize; i++ {
				idx := base + i
				if idx >= TotalKeys { continue }
				// Small value
				val := "v"
				if err := e.Set(keys[idx], val, uint64(idx)); err != nil {
					panic(err)
				}
				atomic.AddInt64(&insertOps, 1)
			}
		}(w)
	}
	wg.Wait()
	
	duration := time.Since(start)
	qps := float64(TotalKeys) / duration.Seconds()
	currentMem := getMemUsage()
	printStats("Insert", TotalKeys, duration, qps, getFileSize(DataDir+"/values.data"), currentMem - startMem)

	// 2. Query Consistency Verification
	fmt.Println("\n--- Phase 2: Query Consistency Verification ---")
	
	verifyQuery(e, "Exact Match (Len 1)", `key = "root0"`, 1)
	verifyQuery(e, "Exact Match (Len 2)", `key = "user.0"`, 1)
	verifyQuery(e, "Prefix Scan (Len 3)", `key like "group.0.member.*"`, 200) // 20000 items / 5 types = 4000 items. id%100 -> 100 groups. 4000/100 = 40? Wait.
	// Total items in group 2: 20,000.
	// id goes from 0 to 19999.
	// id%100 goes 0..99.
	// So each group (0..99) appears 20000/100 = 200 times.
	// So "group.0.member.*" should have 200 matches. Correct.
	
	verifyQuery(e, "Prefix Scan (Len 4)", `key like "app.0.ver.0.config.*"`, 2000)
	// Group 3 items: 20000.
	// Key format: app.<id%10>.ver.<id%10>.config.<id>
	// Condition: id%10 == 0.
	// Since id goes 0..19999, exactly 1/10 of ids satisfy (id%10 == 0).
	// Count = 20000 / 10 = 2000.
	
	verifyQuery(e, "Prefix Scan (Len 5)", `key like "log.2023.*"`, 20000) // All items in group 4

	// 3. Query Performance Test (Prefix)
	fmt.Println("\n--- Phase 3: Query Performance (Prefix Scan) ---")
	
	start = time.Now()
	qCount := 10000
	wg = sync.WaitGroup{}
	chunkSize = qCount / workers
	
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < chunkSize; i++ {
				// Query overlapping prefixes to stress test binary search + scan
				// "group.50.member.*" -> 200 items scan
				_, _ = e.Query(`key like "group.50.member.*"`)
			}
		}()
	}
	wg.Wait()
	duration = time.Since(start)
	qps = float64(qCount) / duration.Seconds()
	printStats("Query(Prefix)", qCount, duration, qps, 0, 0)

	// Final Memory Report
	runtime.GC()
	finalMem := getMemUsage()
	fmt.Printf("\n[Final Stats] Total Keys: %d | Memory Usage: %.2f MB\n", TotalKeys, float64(finalMem-startMem)/1024/1024)
}

func verifyQuery(e *db.Engine, name, sql string, expected int) {
	start := time.Now()
	res, err := e.Query(sql)
	if err != nil {
		fmt.Printf("FAIL: %s - Error: %v\n", name, err)
		return
	}
	duration := time.Since(start)
	
	if len(res) == expected {
		fmt.Printf("PASS: %-25s | Count: %5d | Time: %v\n", name, len(res), duration)
	} else {
		fmt.Printf("FAIL: %-25s | Expected: %d, Got: %d\n", name, expected, len(res))
	}
}

func randomString(n int) string {
	const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	b := make([]byte, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

func getFileSize(path string) int64 {
	fi, err := os.Stat(path)
	if err != nil {
		return 0
	}
	return fi.Size()
}

func getMemUsage() uint64 {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	return m.Alloc
}

func printStats(op string, count int, d time.Duration, qps float64, size int64, memDelta uint64) {
	sizeStr := ""
	if size > 0 {
		sizeStr = fmt.Sprintf(" | Disk: %.2f MB", float64(size)/1024/1024)
	}
	memStr := ""
	if memDelta > 0 {
		memStr = fmt.Sprintf(" | Mem+: %.2f MB", float64(memDelta)/1024/1024)
	}
	fmt.Printf("%-15s: %6d ops in %7v | QPS: %8.0f%s%s\n", op, count, d, qps, sizeStr, memStr)
}