// Full Index Test for CyborgDB Go SDK
//
// This script tests the service's ability to handle large-scale vector indexing
// using the wiki_all_1M dataset from ann-benchmarks:
// 1. Downloads the wiki_all_1M dataset (if not cached)
// 2. Creates an index
// 3. Inserts vectors in batches
// 4. Trains the index
// 5. Runs queries and measures performance
// 6. Measures recall against precomputed ground truth
// 7. Reports throughput and quality metrics
//
// Usage: go run full_index_test_go.go

package main

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"time"

	"gonum.org/v1/hdf5"

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

const (
	datasetURL      = "https://wiki-all.s3.us-east-1.amazonaws.com/wiki_all_1M.hdf5"
	datasetFilename = "wiki_all_1M.hdf5"
)

func getEnvInt(key string, defaultVal int) int {
	if val := os.Getenv(key); val != "" {
		if i, err := strconv.Atoi(val); err == nil {
			return i
		}
	}
	return defaultVal
}

func getEnvFloat(key string, defaultVal float64) float64 {
	if val := os.Getenv(key); val != "" {
		if f, err := strconv.ParseFloat(val, 64); err == nil {
			return f
		}
	}
	return defaultVal
}

func percentile(data []float64, p float64) float64 {
	if len(data) == 0 {
		return 0
	}
	sorted := make([]float64, len(data))
	copy(sorted, data)
	sort.Float64s(sorted)
	idx := int(float64(len(sorted)-1) * p / 100)
	return sorted[idx]
}

// downloadDataset downloads the wiki_all_1M dataset if not cached
func downloadDataset(cacheDir string) (string, error) {
	if err := os.MkdirAll(cacheDir, 0755); err != nil {
		return "", fmt.Errorf("failed to create cache dir: %w", err)
	}

	datasetPath := filepath.Join(cacheDir, datasetFilename)

	if _, err := os.Stat(datasetPath); err == nil {
		fmt.Printf("  Dataset already cached at %s\n", datasetPath)
		return datasetPath, nil
	}

	fmt.Printf("  Downloading from %s...\n", datasetURL)

	resp, err := http.Get(datasetURL)
	if err != nil {
		return "", fmt.Errorf("failed to download: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("download failed with status: %d", resp.StatusCode)
	}

	out, err := os.Create(datasetPath)
	if err != nil {
		return "", fmt.Errorf("failed to create file: %w", err)
	}
	defer out.Close()

	written, err := io.Copy(out, resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to write file: %w", err)
	}

	fmt.Printf("  Downloaded %d bytes to %s\n", written, datasetPath)
	return datasetPath, nil
}

// loadDataset loads the wiki_all_1M dataset from HDF5 file
func loadDataset(datasetPath string, numVectors, numQueries, topK int) ([][]float32, [][]float32, [][]int32, error) {
	file, err := hdf5.OpenFile(datasetPath, hdf5.F_ACC_RDONLY)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to open HDF5 file: %w", err)
	}
	defer file.Close()

	// Load train vectors
	trainDataset, err := file.OpenDataset("train")
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to open train dataset: %w", err)
	}
	defer trainDataset.Close()

	trainSpace := trainDataset.Space()
	trainDims, _, _ := trainSpace.SimpleExtentDims()
	totalVectors := int(trainDims[0])
	vectorDim := int(trainDims[1])

	if numVectors > totalVectors {
		numVectors = totalVectors
	}

	// Read the full train dataset (gonum/hdf5 doesn't support partial reads easily)
	trainFlat := make([]float32, totalVectors*vectorDim)
	if err := trainDataset.Read(&trainFlat); err != nil {
		return nil, nil, nil, fmt.Errorf("failed to read train data: %w", err)
	}

	// Slice to requested number of vectors
	trainVectors := make([][]float32, numVectors)
	for i := 0; i < numVectors; i++ {
		trainVectors[i] = trainFlat[i*vectorDim : (i+1)*vectorDim]
	}

	// Load test vectors (queries)
	testDataset, err := file.OpenDataset("test")
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to open test dataset: %w", err)
	}
	defer testDataset.Close()

	testSpace := testDataset.Space()
	testDims, _, _ := testSpace.SimpleExtentDims()
	totalQueries := int(testDims[0])

	if numQueries > totalQueries {
		numQueries = totalQueries
	}

	// Read the full test dataset, then slice
	testFlat := make([]float32, totalQueries*vectorDim)
	if err := testDataset.Read(&testFlat); err != nil {
		return nil, nil, nil, fmt.Errorf("failed to read test data: %w", err)
	}

	queryVectors := make([][]float32, numQueries)
	for i := 0; i < numQueries; i++ {
		queryVectors[i] = testFlat[i*vectorDim : (i+1)*vectorDim]
	}

	// Load ground truth neighbors
	neighborsDataset, err := file.OpenDataset("neighbors")
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to open neighbors dataset: %w", err)
	}
	defer neighborsDataset.Close()

	neighborsSpace := neighborsDataset.Space()
	neighborsDims, _, _ := neighborsSpace.SimpleExtentDims()
	totalNeighborsQueries := int(neighborsDims[0])
	neighborsK := int(neighborsDims[1])

	if topK > neighborsK {
		topK = neighborsK
	}

	// Read the full neighbors dataset, then slice
	neighborsFlat := make([]int32, totalNeighborsQueries*neighborsK)
	if err := neighborsDataset.Read(&neighborsFlat); err != nil {
		return nil, nil, nil, fmt.Errorf("failed to read neighbors data: %w", err)
	}

	groundTruth := make([][]int32, numQueries)
	for i := 0; i < numQueries; i++ {
		groundTruth[i] = neighborsFlat[i*neighborsK : i*neighborsK+topK]
	}

	return trainVectors, queryVectors, groundTruth, nil
}

// computeRecall computes the recall between ground truth and ANN results
func computeRecall(groundTruth []int32, resultIDs []string) float64 {
	if len(groundTruth) == 0 {
		return 0.0
	}

	gtSet := make(map[int32]bool)
	for _, id := range groundTruth {
		gtSet[id] = true
	}

	matches := 0
	for _, idStr := range resultIDs {
		id, err := strconv.ParseInt(idStr, 10, 32)
		if err == nil && gtSet[int32(id)] {
			matches++
		}
	}

	return float64(matches) / float64(len(groundTruth))
}

func main() {
	// Configuration from environment
	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 environment variable not set")
		os.Exit(1)
	}

	vectorCount := getEnvInt("VECTOR_COUNT", 1000000)
	batchSize := getEnvInt("BATCH_SIZE", 10000)
	topK := getEnvInt("TOP_K", 100)
	numQueries := getEnvInt("NUM_QUERIES", 1000)
	recallBaseline := getEnvFloat("RECALL_BASELINE", 0.95)
	cacheDir := os.Getenv("DATASET_CACHE_DIR")
	if cacheDir == "" {
		cacheDir = "/tmp/ann-benchmarks"
	}

	fmt.Println("============================================================")
	fmt.Println("CyborgDB Full Index Test - Go SDK (wiki_all_1M)")
	fmt.Println("============================================================")
	fmt.Printf("Service URL: %s\n", baseURL)
	fmt.Printf("Vector count: %d\n", vectorCount)
	fmt.Printf("Batch size: %d\n", batchSize)
	fmt.Printf("Top K: %d\n", topK)
	fmt.Printf("Number of queries: %d\n", numQueries)
	fmt.Printf("Recall baseline: %.0f%%\n", recallBaseline*100)
	fmt.Println("============================================================")

	ctx := context.Background()

	// Download dataset
	fmt.Println("\n[1/8] Downloading/loading wiki_all_1M dataset...")
	datasetPath, err := downloadDataset(cacheDir)
	if err != nil {
		fmt.Printf("ERROR: Failed to download dataset: %v\n", err)
		os.Exit(1)
	}

	// Load dataset
	fmt.Println("\n[2/8] Loading dataset...")
	trainVectors, queryVectors, groundTruth, err := loadDataset(datasetPath, vectorCount, numQueries, topK)
	if err != nil {
		fmt.Printf("ERROR: Failed to load dataset: %v\n", err)
		os.Exit(1)
	}

	vectorDim := len(trainVectors[0])
	actualVectorCount := len(trainVectors)

	fmt.Printf("  Train vectors: %d x %d\n", actualVectorCount, vectorDim)
	fmt.Printf("  Query vectors: %d x %d\n", len(queryVectors), vectorDim)
	fmt.Printf("  Ground truth: %d x %d\n", len(groundTruth), len(groundTruth[0]))

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

	// Check health
	health, err := client.GetHealth(ctx)
	if err != nil {
		fmt.Printf("WARNING: Health check failed: %v\n", err)
	} else {
		fmt.Printf("  Service health: %s\n", health["status"])
	}

	// Generate encryption key
	indexKey, err := cyborgdb.GenerateKey()
	if err != nil {
		fmt.Printf("ERROR: Failed to generate key: %v\n", err)
		os.Exit(1)
	}

	// Create index
	indexName := fmt.Sprintf("wiki_full_index_test_%d", time.Now().Unix())
	fmt.Printf("\n[4/8] Creating 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)
		os.Exit(1)
	}
	fmt.Println("  Index created successfully")

	// Track if test passed for exit code
	testPassed := true

	// Ensure cleanup
	defer func() {
		fmt.Println("\n[8/8] Cleaning up...")
		if err := index.DeleteIndex(ctx); err != nil {
			fmt.Printf("  Warning: Failed to delete index: %v\n", err)
		} else {
			fmt.Printf("  Index '%s' deleted\n", indexName)
		}

		if !testPassed {
			os.Exit(1)
		}
	}()

	// Insert vectors in batches
	fmt.Printf("\n[5/8] Inserting %d vectors...\n", actualVectorCount)
	insertStart := time.Now()
	totalInserted := 0

	for batchStart := 0; batchStart < actualVectorCount; batchStart += batchSize {
		batchEnd := batchStart + batchSize
		if batchEnd > actualVectorCount {
			batchEnd = actualVectorCount
		}
		currentBatchSize := batchEnd - batchStart

		items := make(cyborgdb.VectorItems, currentBatchSize)
		for i := 0; i < currentBatchSize; i++ {
			items[i] = cyborgdb.VectorItem{
				Id:     strconv.Itoa(batchStart + i),
				Vector: trainVectors[batchStart+i],
			}
		}

		if err := index.Upsert(ctx, items); err != nil {
			fmt.Printf("ERROR: Upsert failed at batch %d: %v\n", batchStart, err)
			os.Exit(1)
		}

		totalInserted += currentBatchSize
		progress := float64(totalInserted) / float64(actualVectorCount) * 100
		fmt.Printf("\r  Progress: %d/%d (%.1f%%)", totalInserted, actualVectorCount, progress)
	}

	insertTime := time.Since(insertStart).Seconds()
	insertRate := float64(totalInserted) / insertTime

	fmt.Printf("\n\nInsert Results:\n")
	fmt.Printf("  Total vectors: %d\n", totalInserted)
	fmt.Printf("  Total time: %.2fs\n", insertTime)
	fmt.Printf("  Throughput: %.2f vectors/s\n", insertRate)

	// Train the index (auto training)
	fmt.Println("\n[6/8] Training index (auto)...")
	trainStart := time.Now()
	nLists := int32(0)
	if err := index.Train(ctx, cyborgdb.TrainParams{NLists: &nLists}); err != nil {
		fmt.Printf("ERROR: Training failed: %v\n", err)
		os.Exit(1)
	}
	trainTime := time.Since(trainStart).Seconds()
	fmt.Printf("  Training completed in %.2fs\n", trainTime)

	// Use auto n_probes (nil) - the service will calculate optimal value
	fmt.Println("  Using n_probes: auto")

	// Run queries and measure recall
	fmt.Printf("\n[7/8] Running %d queries...\n", numQueries)
	queryTimes := make([]float64, numQueries)
	recalls := make([]float64, 0, numQueries)

	for i := 0; i < numQueries; i++ {
		queryVector := queryVectors[i]

		queryStart := time.Now()
		results, err := index.Query(ctx, cyborgdb.QueryParams{
			QueryVector: queryVector,
			TopK:        int32(topK),
			// NProbes: nil for auto
		})
		queryTime := time.Since(queryStart).Seconds()
		queryTimes[i] = queryTime

		if err != nil {
			fmt.Printf("WARNING: Query %d failed: %v\n", i, err)
			continue
		}

		// Extract result IDs and compute recall
		resultsData := results.GetResults()
		if resultsData.ArrayOfQueryResultItem != nil && len(*resultsData.ArrayOfQueryResultItem) > 0 {
			resultIDs := make([]string, 0, len(*resultsData.ArrayOfQueryResultItem))
			for _, item := range *resultsData.ArrayOfQueryResultItem {
				resultIDs = append(resultIDs, item.Id)
			}

			recall := computeRecall(groundTruth[i], resultIDs)
			recalls = append(recalls, recall)

			if i == 0 {
				fmt.Printf("  First query returned %d results\n", len(*resultsData.ArrayOfQueryResultItem))
			}
		} else {
			if i == 0 {
				fmt.Println("  WARNING: First query returned no results")
			}
			recalls = append(recalls, 0.0)
		}

		if (i+1)%100 == 0 {
			fmt.Printf("\r  Progress: %d/%d", i+1, numQueries)
		}
	}

	// Calculate query statistics
	var sum float64
	for _, t := range queryTimes {
		sum += t
	}
	avgQueryTime := sum / float64(numQueries)
	p50QueryTime := percentile(queryTimes, 50)
	p95QueryTime := percentile(queryTimes, 95)
	p99QueryTime := percentile(queryTimes, 99)

	fmt.Printf("\n\nQuery Results:\n")
	fmt.Printf("  Total queries: %d\n", numQueries)
	fmt.Printf("  Average latency: %.2fms\n", avgQueryTime*1000)
	fmt.Printf("  P50 latency: %.2fms\n", p50QueryTime*1000)
	fmt.Printf("  P95 latency: %.2fms\n", p95QueryTime*1000)
	fmt.Printf("  P99 latency: %.2fms\n", p99QueryTime*1000)
	fmt.Printf("  QPS: %.2f\n", 1/avgQueryTime)

	// Calculate recall statistics
	var avgRecall, minRecall, maxRecall float64
	if len(recalls) > 0 {
		var recallSum float64
		minRecall = recalls[0]
		maxRecall = recalls[0]
		for _, r := range recalls {
			recallSum += r
			if r < minRecall {
				minRecall = r
			}
			if r > maxRecall {
				maxRecall = r
			}
		}
		avgRecall = recallSum / float64(len(recalls))

		fmt.Printf("\nRecall Results (Recall@%d):\n", topK)
		fmt.Printf("  Average recall: %.2f%%\n", avgRecall*100)
		fmt.Printf("  Min recall: %.2f%%\n", minRecall*100)
		fmt.Printf("  Max recall: %.2f%%\n", maxRecall*100)
		fmt.Printf("  Baseline threshold: %.2f%%\n", recallBaseline*100)

		if avgRecall >= recallBaseline {
			fmt.Printf("  ✓ PASS: Recall %.2f%% >= baseline %.2f%%\n", avgRecall*100, recallBaseline*100)
		} else {
			fmt.Printf("  ✗ FAIL: Recall %.2f%% < baseline %.2f%%\n", avgRecall*100, recallBaseline*100)
			testPassed = false
		}
	}

	// Summary
	fmt.Println("\n============================================================")
	fmt.Println("FULL INDEX TEST COMPLETED")
	fmt.Println("============================================================")
	fmt.Printf("Dataset: wiki_all_1M\n")
	fmt.Printf("Index: %s\n", indexName)
	fmt.Printf("Vectors indexed: %d\n", totalInserted)
	fmt.Printf("Vector dimension: %d\n", vectorDim)
	fmt.Printf("Insert throughput: %.2f vectors/s\n", insertRate)
	fmt.Printf("Training time: %.2fs\n", trainTime)
	fmt.Printf("Query latency (avg): %.2fms\n", avgQueryTime*1000)
	if len(recalls) > 0 {
		fmt.Printf("Recall@%d: %.2f%%\n", topK, avgRecall*100)
	}
	fmt.Println("============================================================")

	if testPassed {
		fmt.Println("\nTest completed successfully!")
	} else {
		fmt.Println("\nTest FAILED: Recall below baseline threshold!")
	}
}
