package main

import (
	"context"
	"fmt"
	"time"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

func main() {
	// Connect to NATS
	nc, err := nats.Connect(nats.DefaultURL)
	if err != nil {
		panic(err)
	}
	defer nc.Close()
	// Create JetStream context
	js, err := jetstream.New(nc)
	if err != nil {
		panic(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	streamName := "GRID_test-grid"
	bucket := "test-grid"
	// Delete stream if exists (this also deletes all consumers)
	js.DeleteStream(ctx, streamName)
	// Create stream matching Lua configuration
	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
		Name:        streamName,
		Subjects:    []string{"grid." + bucket, "change." + bucket},
		Storage:     jetstream.FileStorage,
		Replicas:    1,
		MaxAge:      24 * time.Hour,
		Retention:   jetstream.LimitsPolicy,
		Discard:     jetstream.DiscardNew,
		AllowDirect: true,
	})
	if err != nil {
		fmt.Printf("Stream create error: %v\n", err)
	}
	// Publish test messages matching Lua grid format
	fmt.Println("Publishing 50 test messages...")
	for i := 1; i <= 50; i++ {
		msg := fmt.Sprintf(`{"message_type":"grid","message_id":"msg%d","bucket":"%s","chunk_index":0,"total_chunks":1,"row_count":1,"total_rows":50,"data":{"record_id":"id%d","field1":"value%d"}}`, i, bucket, i, i)
		_, err = js.Publish(ctx, "grid."+bucket, []byte(msg))
		if err != nil {
			fmt.Printf("Publish error: %v\n", err)
		}
	}
	fmt.Printf("Published 50 messages to %s\n", "grid."+bucket)
	// Get stream info
	stream, err := js.Stream(ctx, streamName)
	if err != nil {
		panic(err)
	}
	info, err := stream.Info(ctx)
	if err != nil {
		panic(err)
	}
	fmt.Printf("Stream state: %d messages\n", info.State.Msgs)
	// Test 1: Fetch(50) - exact match for available messages
	fmt.Println("\n=== Test 1: Fetch(50) with ord_1 ===")
	cons1, err := stream.CreateConsumer(ctx, jetstream.ConsumerConfig{
		Durable:           "ord_1",
		AckPolicy:         jetstream.AckNonePolicy,
		DeliverPolicy:     jetstream.DeliverAllPolicy,
		ReplayPolicy:      jetstream.ReplayInstantPolicy,
		MemoryStorage:     true,
		FilterSubject:     "grid." + bucket,
		InactiveThreshold: 5 * time.Minute,
	})
	if err != nil {
		panic(err)
	}
	fmt.Printf("Created consumer: %s\n", cons1.CachedInfo().Name)
	startTime := time.Now()
	batch1, err := cons1.Fetch(50,
		jetstream.FetchMaxWait(2*time.Second),
	)
	if err != nil {
		fmt.Printf("Fetch error: %v\n", err)
	} else {
		msgCount := 0
		for msg := range batch1.Messages() {
			msgCount++
			if msgCount <= 3 {
				fmt.Printf("  Message %d: %s\n", msgCount, string(msg.Data()))
			}
		}
		elapsed := time.Since(startTime)
		fmt.Printf("Fetched %d messages in %v\n", msgCount, elapsed)
		if batch1.Error() != nil {
			fmt.Printf("Batch error: %v\n", batch1.Error())
		}
	}
	// Test 2: Fetch(50) with ANOTHER fresh consumer
	fmt.Println("\n=== Test 2: Fetch(50) with ord_2 ===")
	cons2, err := stream.CreateConsumer(ctx, jetstream.ConsumerConfig{
		Durable:           "ord_2",
		AckPolicy:         jetstream.AckNonePolicy,
		DeliverPolicy:     jetstream.DeliverAllPolicy,
		ReplayPolicy:      jetstream.ReplayInstantPolicy,
		MemoryStorage:     true,
		FilterSubject:     "grid." + bucket,
		InactiveThreshold: 5 * time.Minute,
	})
	if err != nil {
		panic(err)
	}
	fmt.Printf("Created consumer: %s\n", cons2.CachedInfo().Name)
	startTime = time.Now()
	batch2, err := cons2.Fetch(50,
		jetstream.FetchMaxWait(2*time.Second),
	)
	if err != nil {
		fmt.Printf("Fetch error: %v\n", err)
	} else {
		msgCount := 0
		for msg := range batch2.Messages() {
			msgCount++
			if msgCount <= 3 {
				fmt.Printf("  Message %d: %s\n", msgCount, string(msg.Data()))
			}
		}
		elapsed := time.Since(startTime)
		fmt.Printf("Fetched %d messages in %v\n", msgCount, elapsed)
		if batch2.Error() != nil {
			fmt.Printf("Batch error: %v\n", batch2.Error())
		}
	}
	// Cleanup
	js.DeleteStream(ctx, streamName)
	fmt.Println("\nStream deleted - test complete")
}
