package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"

	deltalake "github.com/tidbcloud/metering/delta-rs-golang"
)

func main() {
	// Create a temporary directory for the Delta table
	tempDir, err := os.MkdirTemp("/tmp", "delta_partitioned_table")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tempDir) // Clean up when done

	tablePath := filepath.Join(tempDir, "partitioned_table")

	// Define the table schema with columns including partition columns
	columns := []deltalake.ColumnDefinition{
		{Name: "id", Type: "integer", Nullable: false},
		{Name: "name", Type: "string", Nullable: true},
		{Name: "age", Type: "integer", Nullable: true},
		{Name: "year", Type: "integer", Nullable: false},
		{Name: "month", Type: "integer", Nullable: false},
	}

	// Define partition columns
	partitionColumns := []string{"year", "month"}

	fmt.Printf("Creating table with partition columns: %v\n", partitionColumns)

	// Create table options with partition columns
	options := deltalake.CreateTableOptions{
		TablePath:        tablePath,
		SaveMode:         deltalake.WriteModeOverwrite,
		Columns:          columns,
		PartitionColumns: partitionColumns,
		TableName:        stringPtr("users_partitioned"),
		StorageOptions:   nil,
	}

	// Create the partitioned Delta table
	table, err := deltalake.CreateDeltaLakeTable(options)
	if err != nil {
		log.Fatalf("Failed to create partitioned table: %v", err)
	}
	defer table.Close()

	fmt.Printf("Successfully created partitioned Delta table at: %s\n", tablePath)

	// Get table version
	version, err := table.Version()
	if err != nil {
		log.Fatalf("Failed to get table version: %v", err)
	}
	fmt.Printf("Table version: %d\n", version)

	// Get table files
	files, err := table.GetFiles()
	if err != nil {
		log.Fatalf("Failed to get table files: %v", err)
	}
	fmt.Printf("Table files: %v\n", files)

	fmt.Println("Partitioned table example completed successfully!")
}

// Helper function to get a string pointer
func stringPtr(s string) *string {
	return &s
}