package main

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

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

func main() {
	// Create a temporary directory for the demo
	tempDir, err := os.MkdirTemp("/tmp", "delta_write_example")
	if err != nil {
		log.Fatal(err)
	}
	// defer os.RemoveAll(tempDir) // Clean up when done - commented out for inspection

	tablePath := filepath.Join(tempDir, "employee_table")
	fmt.Printf("Creating Delta Lake table at: %s\n", tablePath)

	// Step 1: Create a new Delta table
	fmt.Println("\n📊 Step 1: Creating Delta table...")

	// Define the table schema
	columns := []deltalake.ColumnDefinition{
		{Name: "id", Type: "integer", Nullable: false},
		{Name: "name", Type: "string", Nullable: true},
		{Name: "age", Type: "integer", Nullable: true},
		{Name: "department", Type: "string", Nullable: true},
		{Name: "salary", Type: "double", Nullable: true},
	}

	// Create table options
	tableName := "employee_table"
	options := deltalake.CreateTableOptions{
		TablePath: tablePath,
		SaveMode:  deltalake.WriteModeOverwrite,
		Columns:   columns,
		TableName: &tableName,
	}

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

	fmt.Println("✅ Delta Lake table created successfully!")

	// Get initial table info
	version, err := table.Version()
	if err != nil {
		log.Fatal("Failed to get table version:", err)
	}
	fmt.Printf("📋 Initial table version: %d\n", version)

	// Step 2: Write first batch of data
	fmt.Println("\n✍️  Step 2: Writing first batch of employee data...")

	writeData1 := deltalake.WriteData{
		Schema: columns,
		Data: [][]interface{}{
			{1, "Alice Johnson", 28, "Engineering", 75000.0},
			{2, "Bob Smith", 35, "Marketing", 65000.0},
			{3, "Charlie Brown", 42, "Engineering", 85000.0},
		},
	}

	// Write data to the table using Append mode
	updatedTable1, err := table.Write(writeData1, deltalake.WriteModeAppend)
	if err != nil {
		log.Fatal("Failed to write first batch:", err)
	}
	defer updatedTable1.Close()

	fmt.Println("✅ First batch written successfully!")

	// Get updated table info
	version1, err := updatedTable1.Version()
	if err != nil {
		log.Fatal("Failed to get table version:", err)
	}
	fmt.Printf("📋 Table version after first write: %d\n", version1)

	// Get table files
	files1, err := updatedTable1.GetFiles()
	if err != nil {
		log.Fatal("Failed to get table files:", err)
	}
	fmt.Printf("📁 Number of files after first write: %d\n", len(files1))

	// Step 3: Write second batch of data
	fmt.Println("\n✍️  Step 3: Writing second batch of employee data...")

	writeData2 := deltalake.WriteData{
		Schema: columns,
		Data: [][]interface{}{
			{4, "Diana Prince", 30, "HR", 60000.0},
			{5, "Edward Wilson", 25, "Engineering", 70000.0},
			{6, "Fiona Green", 38, "Finance", 80000.0},
		},
	}

	// Write second batch
	updatedTable2, err := updatedTable1.Write(writeData2, deltalake.WriteModeAppend)
	if err != nil {
		log.Fatal("Failed to write second batch:", err)
	}
	defer updatedTable2.Close()

	fmt.Println("✅ Second batch written successfully!")

	// Get final table info
	version2, err := updatedTable2.Version()
	if err != nil {
		log.Fatal("Failed to get table version:", err)
	}
	fmt.Printf("📋 Final table version: %d\n", version2)

	files2, err := updatedTable2.GetFiles()
	if err != nil {
		log.Fatal("Failed to get table files:", err)
	}
	fmt.Printf("📁 Total number of files: %d\n", len(files2))

	// Step 4: Demonstrate different data types
	fmt.Println("\n🔢 Step 4: Writing data with different types...")

	// Create a table with various data types
	diverseTablePath := filepath.Join(tempDir, "diverse_table")
	diverseColumns := []deltalake.ColumnDefinition{
		{Name: "int_col", Type: "integer", Nullable: false},
		{Name: "long_col", Type: "long", Nullable: true},
		{Name: "string_col", Type: "string", Nullable: true},
		{Name: "bool_col", Type: "boolean", Nullable: true},
		{Name: "float_col", Type: "float", Nullable: true},
		{Name: "double_col", Type: "double", Nullable: true},
	}

	diverseTableName := "diverse_table"
	diverseOptions := deltalake.CreateTableOptions{
		TablePath: diverseTablePath,
		SaveMode:  deltalake.WriteModeOverwrite,
		Columns:   diverseColumns,
		TableName: &diverseTableName,
	}

	diverseTable, err := deltalake.CreateDeltaLakeTable(diverseOptions)
	if err != nil {
		log.Fatal("Failed to create diverse table:", err)
	}
	defer diverseTable.Close()

	diverseData := deltalake.WriteData{
		Schema: diverseColumns,
		Data: [][]interface{}{
			{123, int64(9876543210), "Hello World", true, float32(3.14), 2.718281828},
			{456, int64(1234567890), "Delta Lake", false, float32(1.41), 1.414213562},
			{789, int64(5555555555), "Go Programming", true, float32(2.71), 3.141592654},
		},
	}

	diverseUpdated, err := diverseTable.Write(diverseData, deltalake.WriteModeAppend)
	if err != nil {
		log.Fatal("Failed to write diverse data:", err)
	}
	defer diverseUpdated.Close()

	fmt.Println("✅ Diverse data types written successfully!")

	diverseVersion, err := diverseUpdated.Version()
	if err != nil {
		log.Fatal("Failed to get diverse table version:", err)
	}
	fmt.Printf("📋 Diverse table version: %d\n", diverseVersion)

	// Final summary
	fmt.Println("\n🎉 Write Data Example Completed Successfully!")
	fmt.Println("Summary:")
	fmt.Printf("  • Created 2 Delta Lake tables\n")
	fmt.Printf("  • Employee table has %d records across %d files (version %d)\n", 6, len(files2), version2)
	fmt.Printf("  • Diverse table has 3 records with various data types (version %d)\n", diverseVersion)
	fmt.Printf("  • Tables stored in: %s\n", tempDir)

	fmt.Println("\n🔍 You can inspect the created tables using:")
	fmt.Printf("  ls -la %s/*/\n", tempDir)
	fmt.Printf("  ls -la %s/*/_delta_log/\n", tempDir)
}