// Persistence Test for CyborgDB Go SDK
//
// This script tests data persistence across service restarts:
// - Phase 1 (setup): Creates index, inserts data, saves index info to file
// - Phase 2 (verify): Loads index, verifies data still exists
//
// Usage:
//   go run persistence_test_go.go setup    # Run before restart
//   go run persistence_test_go.go verify   # Run after restart

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"time"

	cyborgdb "github.com/cyborginc/cyborgdb-go"
)

const (
	stateFile   = "/tmp/persistence_test_state_go.json"
	vectorCount = 1000
	vectorDim   = 128
)

type TestState struct {
	IndexName   string   `json:"index_name"`
	IndexKey    string   `json:"index_key"`
	VectorCount int      `json:"vector_count"`
	VectorDim   int      `json:"vector_dim"`
	TestIDs     []string `json:"test_ids"`
}

func generateSeededVector(dim int, seed int64) []float32 {
	r := rand.New(rand.NewSource(seed))
	vec := make([]float32, dim)
	for i := range vec {
		vec[i] = float32(r.NormFloat64())
	}
	return vec
}

func setupPhase() int {
	fmt.Println("============================================================")
	fmt.Println("Persistence Test - Setup Phase")
	fmt.Println("============================================================")

	ctx := context.Background()
	baseURL := os.Getenv("CYBORGDB_URL")
	if baseURL == "" {
		baseURL = "http://localhost:8000"
	}
	apiKey := os.Getenv("CYBORGDB_API_KEY")
	if apiKey == "" {
		fmt.Println("ERROR: CYBORGDB_API_KEY not set")
		return 1
	}

	// Initialize client
	fmt.Println("\n[1/4] Initializing client...")
	client, err := cyborgdb.NewClient(baseURL, apiKey)
	if err != nil {
		fmt.Printf("ERROR: Failed to create client: %v\n", err)
		return 1
	}

	// Generate key and create index
	fmt.Println("[2/4] Creating index...")
	indexKey, err := cyborgdb.GenerateKey()
	if err != nil {
		fmt.Printf("ERROR: Failed to generate key: %v\n", err)
		return 1
	}

	indexName := "persistence_test_index_go"

	// Try to delete existing index
	if existing, err := client.LoadIndex(ctx, indexName, indexKey); err == nil {
		existing.DeleteIndex(ctx)
		fmt.Printf("  Deleted existing index '%s'\n", indexName)
	}

	indexConfig := cyborgdb.IndexIVFFlat(int32(vectorDim))
	metric := "euclidean"

	index, err := client.CreateIndex(ctx, &cyborgdb.CreateIndexParams{
		IndexName:   indexName,
		IndexKey:    indexKey,
		IndexConfig: indexConfig,
		Metric:      &metric,
	})
	if err != nil {
		fmt.Printf("ERROR: Failed to create index: %v\n", err)
		return 1
	}
	fmt.Printf("  Created index '%s'\n", indexName)

	// Insert test vectors with seeded random values
	fmt.Printf("[3/4] Inserting %d test vectors...\n", vectorCount)

	items := make(cyborgdb.VectorItems, vectorCount)
	testIDs := make([]string, 10)

	for i := 0; i < vectorCount; i++ {
		id := fmt.Sprintf("persist_vec_%d", i)
		items[i] = cyborgdb.VectorItem{
			Id:     id,
			Vector: generateSeededVector(vectorDim, int64(42+i)),
		}
		if i < 10 {
			testIDs[i] = id
		}
	}

	if err := index.Upsert(ctx, items); err != nil {
		fmt.Printf("ERROR: Upsert failed: %v\n", err)
		return 1
	}
	fmt.Printf("  Inserted %d vectors\n", vectorCount)

	// Verify data was inserted
	listResp, err := index.ListIDs(ctx)
	if err != nil {
		fmt.Printf("ERROR: Failed to list IDs: %v\n", err)
		return 1
	}
	fmt.Printf("  Verified %d vectors in index\n", len(listResp.Ids))

	// Save state for verification phase
	fmt.Println("[4/4] Saving state...")
	state := TestState{
		IndexName:   indexName,
		IndexKey:    base64.StdEncoding.EncodeToString(indexKey),
		VectorCount: vectorCount,
		VectorDim:   vectorDim,
		TestIDs:     testIDs,
	}

	stateBytes, _ := json.Marshal(state)
	if err := os.WriteFile(stateFile, stateBytes, 0644); err != nil {
		fmt.Printf("ERROR: Failed to save state: %v\n", err)
		return 1
	}
	fmt.Printf("  State saved to %s\n", stateFile)

	fmt.Println("\n============================================================")
	fmt.Println("Setup phase completed successfully!")
	fmt.Println("Service can now be restarted.")
	fmt.Println("============================================================")
	return 0
}

func verifyPhase() int {
	fmt.Println("============================================================")
	fmt.Println("Persistence Test - Verify Phase")
	fmt.Println("============================================================")

	ctx := context.Background()
	baseURL := os.Getenv("CYBORGDB_URL")
	if baseURL == "" {
		baseURL = "http://localhost:8000"
	}
	apiKey := os.Getenv("CYBORGDB_API_KEY")
	if apiKey == "" {
		fmt.Println("ERROR: CYBORGDB_API_KEY not set")
		return 1
	}

	// Load state from setup phase
	fmt.Println("\n[1/4] Loading state...")
	stateBytes, err := os.ReadFile(stateFile)
	if err != nil {
		fmt.Printf("ERROR: State file not found: %s\n", stateFile)
		fmt.Println("Did you run the setup phase first?")
		return 1
	}

	var state TestState
	if err := json.Unmarshal(stateBytes, &state); err != nil {
		fmt.Printf("ERROR: Failed to parse state: %v\n", err)
		return 1
	}

	indexKey, _ := base64.StdEncoding.DecodeString(state.IndexKey)
	fmt.Printf("  Index: %s\n", state.IndexName)
	fmt.Printf("  Expected vectors: %d\n", state.VectorCount)

	// Initialize client
	fmt.Println("[2/4] Initializing client...")
	client, err := cyborgdb.NewClient(baseURL, apiKey)
	if err != nil {
		fmt.Printf("ERROR: Failed to create client: %v\n", err)
		return 1
	}

	// Load existing index
	fmt.Println("[3/4] Loading index...")
	index, err := client.LoadIndex(ctx, state.IndexName, indexKey)
	if err != nil {
		fmt.Printf("ERROR: Failed to load index: %v\n", err)
		return 1
	}
	fmt.Printf("  Index '%s' loaded successfully\n", state.IndexName)

	// Verify data
	fmt.Println("[4/4] Verifying data...")

	// Check vector count
	listResp, err := index.ListIDs(ctx)
	if err != nil {
		fmt.Printf("ERROR: Failed to list IDs: %v\n", err)
		return 1
	}

	actualCount := len(listResp.Ids)
	fmt.Printf("  Found %d vectors (expected %d)\n", actualCount, state.VectorCount)

	if actualCount != state.VectorCount {
		fmt.Println("ERROR: Vector count mismatch!")
		return 1
	}

	// Verify specific vectors exist
	items, err := index.Get(ctx, state.TestIDs, []string{"vector"})
	if err != nil {
		fmt.Printf("ERROR: Failed to get test vectors: %v\n", err)
		return 1
	}

	if len(items.Results) != len(state.TestIDs) {
		fmt.Printf("ERROR: Expected %d items, got %d\n", len(state.TestIDs), len(items.Results))
		return 1
	}
	fmt.Printf("  Verified %d test vectors exist\n", len(state.TestIDs))

	// Test query still works
	queryVector := generateSeededVector(state.VectorDim, 42)
	results, err := index.Query(ctx, cyborgdb.QueryParams{
		QueryVector: queryVector,
		TopK:        10,
	})
	if err != nil {
		fmt.Printf("ERROR: Query failed: %v\n", err)
		return 1
	}

	resultsData := results.GetResults()
	if resultsData.ArrayOfQueryResultItem == nil || len(*resultsData.ArrayOfQueryResultItem) == 0 {
		fmt.Println("ERROR: Query returned no results")
		return 1
	}
	fmt.Println("  Query returned results successfully")

	// Cleanup
	fmt.Println("\nCleaning up...")
	if err := index.DeleteIndex(ctx); err != nil {
		fmt.Printf("Warning: Failed to delete index: %v\n", err)
	}
	os.Remove(stateFile)
	fmt.Println("  Deleted index and state file")

	fmt.Println("\n============================================================")
	fmt.Println("PERSISTENCE TEST PASSED!")
	fmt.Println("Data successfully persisted across service restart.")
	fmt.Println("============================================================")
	return 0
}

func main() {
	rand.Seed(time.Now().UnixNano())

	if len(os.Args) < 2 {
		fmt.Println("Usage: go run persistence_test_go.go <setup|verify>")
		os.Exit(1)
	}

	phase := os.Args[1]

	var exitCode int
	switch phase {
	case "setup":
		exitCode = setupPhase()
	case "verify":
		exitCode = verifyPhase()
	default:
		fmt.Printf("Unknown phase: %s\n", phase)
		fmt.Println("Usage: go run persistence_test_go.go <setup|verify>")
		exitCode = 1
	}

	os.Exit(exitCode)
}
