package main

import (
	"context"
	"fmt"
	"log"
	"os"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/sts"

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

// AWSCredentials holds the AWS credentials
type AWSCredentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SessionToken    string
	Region          string
}

// getAWSCredentialsFromProfile retrieves AWS credentials from the specified profile
func getAWSCredentialsFromProfile(ctx context.Context, profile string) (*AWSCredentials, error) {
	// Load AWS config with the specified profile
	cfg, err := config.LoadDefaultConfig(ctx,
		config.WithSharedConfigProfile(profile),
		config.WithRegion("us-west-2"), // Default region, can be overridden
	)
	if err != nil {
		return nil, fmt.Errorf("failed to load AWS config: %w", err)
	}

	// Get credentials from the config
	creds, err := cfg.Credentials.Retrieve(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve AWS credentials: %w", err)
	}

	// Verify credentials by calling STS GetCallerIdentity
	stsClient := sts.NewFromConfig(cfg)
	identity, err := stsClient.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
	if err != nil {
		return nil, fmt.Errorf("failed to verify AWS credentials: %w", err)
	}

	fmt.Printf("Successfully retrieved credentials for AWS Account: %s\n", aws.ToString(identity.Account))
	fmt.Printf("User ARN: %s\n", aws.ToString(identity.Arn))

	return &AWSCredentials{
		AccessKeyID:     creds.AccessKeyID,
		SecretAccessKey: creds.SecretAccessKey,
		SessionToken:    creds.SessionToken,
		Region:          cfg.Region,
	}, nil
}

// createStorageOptionsFromCredentials converts AWS credentials to storage options map
func createStorageOptionsFromCredentials(creds *AWSCredentials, bucket string) map[string]string {
	storageOptions := map[string]string{
		"AWS_ACCESS_KEY_ID":     creds.AccessKeyID,
		"AWS_SECRET_ACCESS_KEY": creds.SecretAccessKey,
		"AWS_REGION":            creds.Region,
	}

	// Add session token if present (for temporary credentials)
	if creds.SessionToken != "" {
		storageOptions["AWS_SESSION_TOKEN"] = creds.SessionToken
	}

	// Add S3-specific options for better performance and safety
	storageOptions["AWS_S3_ALLOW_UNSAFE_RENAME"] = "true"

	return storageOptions
}

func main() {
	ctx := context.Background()

	// Get AWS profile from environment variable or use default
	profile := os.Getenv("AWS_PROFILE")
	if profile == "" {
		profile = "ob-dev"
	}

	// Get S3 bucket from environment variable
	bucket := os.Getenv("S3_BUCKET")
	if bucket == "" {
		log.Println("S3_BUCKET environment variable not set, using example bucket")
		bucket = "o11y-dev-shared-us-west-2"
	}

	fmt.Printf("Delta Lake Go Binding - Open Table Example\n")
	fmt.Printf("Using AWS Profile: %s\n", profile)
	fmt.Printf("Target S3 Bucket: %s\n", bucket)

	// Retrieve AWS credentials from the profile
	fmt.Println("\nRetrieving AWS credentials from profile...")
	creds, err := getAWSCredentialsFromProfile(ctx, profile)
	if err != nil {
		log.Fatalf("Failed to get AWS credentials: %v", err)
	}

	fmt.Printf("AWS Region: %s\n", creds.Region)

	// Create storage options from the retrieved credentials
	storageOptions := createStorageOptionsFromCredentials(creds, bucket)

	// S3 table path - same as s3_storage_options example
	tablePath := fmt.Sprintf("s3://%s/jinweitest/delta-tables/products", bucket)

	fmt.Printf("\nOpening Delta table from S3...\n")
	fmt.Printf("Table Path: %s\n", tablePath)

	// Print storage options (hiding sensitive information)
	fmt.Println("\nStorage Options:")
	for key, value := range storageOptions {
		if key == "AWS_SECRET_ACCESS_KEY" || key == "AWS_SESSION_TOKEN" {
			fmt.Printf("  %s: ***hidden***\n", key)
		} else {
			fmt.Printf("  %s: %s\n", key, value)
		}
	}

	// Open the Delta table
	fmt.Println("\nAttempting to open Delta table...")
	table, err := deltalake.OpenDeltaLakeTable(tablePath, storageOptions)
	if err != nil {
		log.Fatalf("Failed to open Delta table: %v", err)
	}
	defer table.Close()

	fmt.Println("✅ Successfully opened Delta table from S3!")

	// Get table version
	fmt.Println("\nGetting 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
	fmt.Println("\nGetting table files...")
	files, err := table.GetFiles()
	if err != nil {
		log.Fatalf("Failed to get table files: %v", err)
	}
	fmt.Printf("Number of files: %d\n", len(files))

	// Show first few files if any exist
	if len(files) > 0 {
		fmt.Println("Sample files:")
		maxFiles := 5
		if len(files) < maxFiles {
			maxFiles = len(files)
		}
		for i := 0; i < maxFiles; i++ {
			fmt.Printf("  - %s\n", files[i])
		}
		if len(files) > 5 {
			fmt.Printf("  ... and %d more files\n", len(files)-5)
		}
	}

	fmt.Println("\n✅ Successfully demonstrated Delta table operations!")
	fmt.Println("📝 Table opened, version retrieved, and files listed from S3")
}