package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"os"
	"time"

	"github.com/pletorco/fluss-go/pkg/fgo"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	address := os.Getenv("FLUSS_BOOTSTRAP")
	if address == "" {
		address = "localhost:9123"
	}
	client, err := fgo.Open(ctx, fgo.WithBootstrapServers(address))
	if err != nil {
		return err
	}
	defer client.Close()

	// Use a separate database so preparation cannot overwrite Quickstart data.
	path := fgo.TablePath{Database: "release_1_0_demo", Table: "customer_profile"}
	table, err := client.GetTable(ctx, path)
	if err != nil {
		return fmt.Errorf("load %s: %w", path, err)
	}
	if err := table.RequirePrimaryKey(); err != nil {
		return err
	}
	if len(table.Schema.PartitionKey) != 0 {
		return fmt.Errorf("this example requires a non-partitioned table")
	}
	buckets, err := client.ResolveTableBuckets(ctx, fgo.PhysicalTablePath{TablePath: path})
	if err != nil {
		return err
	}
	fmt.Printf("Scanning %s: %d buckets\n", path, len(buckets))
	total := 0
	for _, bucket := range buckets {
		count, err := scanBucket(ctx, client, table, bucket)
		if err != nil {
			return fmt.Errorf("scan bucket %d: %w", bucket.BucketID, err)
		}
		total += count
	}
	fmt.Printf("Completed: %d rows\n", total)
	return nil
}

func scanBucket(ctx context.Context, client *fgo.Client, table fgo.Table, bucket fgo.TableBucket) (int, error) {
	scanner, err := client.NewBatchScanner(ctx, table, bucket, fgo.WithBatchSizeBytes(64*1024))
	if err != nil {
		return 0, err
	}
	// Close even when polling or JSON encoding fails.
	defer scanner.Close()
	count := 0
	for {
		batch, err := scanner.Poll(ctx)
		if err != nil {
			return count, err
		}
		done := batch.Done
		// Encode before releasing memory owned by this result.
		data, encodeErr := rowsJSON(table.Schema.Columns, batch.Rows)
		count += len(batch.Rows)
		batch.Release()
		if encodeErr != nil {
			return count, encodeErr
		}
		fmt.Printf("Bucket %d, done=%t\n%s\n", bucket.BucketID, done, data)
		if done {
			return count, nil
		}
	}
}

// Rows use schema order. Add field names before releasing the response memory.
// Keep DECIMAL values as strings to preserve their exact representation.
func rowsJSON(columns []fgo.Column, rows []fgo.Row) ([]byte, error) {
	objects := make([]map[string]any, 0, len(rows))
	for _, row := range rows {
		if len(row) != len(columns) {
			return nil, fmt.Errorf("row has %d values; schema has %d columns", len(row), len(columns))
		}
		object := make(map[string]any, len(columns))
		for i, column := range columns {
			value := row[i]
			if decimal, ok := value.(*big.Rat); ok && decimal != nil {
				if column.LogicalType == nil {
					return nil, fmt.Errorf("column %s has no DECIMAL scale", column.Name)
				}
				value = decimal.FloatString(column.LogicalType.Scale)
			}
			object[column.Name] = value
		}
		objects = append(objects, object)
	}
	return json.MarshalIndent(objects, "", "  ")
}