//go:build script

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	_ "github.com/duckdb/duckdb-go/v2"
)

const (
	account   = "o11yprodeastus2"
	container = "o11y-prod-shared-eastus2"
	relPath   = "deltalake/1372813089189741281/7553f5ff-4163-4424-a0b4-ca774600a4d9/slowlogs"
)

func getAccessToken() (string, error) {
	federatedTokenFile := os.Getenv("AZURE_FEDERATED_TOKEN_FILE")
	if federatedTokenFile == "" {
		return "", fmt.Errorf("AZURE_FEDERATED_TOKEN_FILE is not set")
	}

	federatedToken, err := os.ReadFile(federatedTokenFile)
	if err != nil {
		return "", fmt.Errorf("failed to read federated token file: %w", err)
	}

	clientID := os.Getenv("AZURE_CLIENT_ID")
	tenantID := os.Getenv("AZURE_TENANT_ID")
	if clientID == "" || tenantID == "" {
		return "", fmt.Errorf("AZURE_CLIENT_ID or AZURE_TENANT_ID is not set")
	}

	tokenURL := fmt.Sprintf("https://login.microsoftonline.com/%s/oauth2/v2.0/token", tenantID)

	data := url.Values{}
	data.Set("client_id", clientID)
	data.Set("scope", "https://storage.azure.com/.default")
	data.Set("client_assertion_type", "urn:ietf:params:oauth:client-assertion-type:jwt-bearer")
	data.Set("client_assertion", strings.TrimSpace(string(federatedToken)))
	data.Set("grant_type", "client_credentials")

	req, err := http.NewRequest("POST", tokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return "", fmt.Errorf("failed to create HTTP request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	client := &http.Client{
		Timeout: 30 * time.Second,
	}

	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to send HTTP request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to read response body: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("Azure AD returned error status code %d: %s", resp.StatusCode, string(body))
	}

	var tokenResponse struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int    `json:"expires_in"`
		TokenType   string `json:"token_type"`
	}

	if err := json.Unmarshal(body, &tokenResponse); err != nil {
		return "", fmt.Errorf("failed to parse token response: %w", err)
	}

	if tokenResponse.AccessToken == "" {
		return "", fmt.Errorf("access token not found in response")
	}

	fmt.Printf("Successfully obtained Azure access token\n")
	return tokenResponse.AccessToken, nil
}

func main() {
	// Get access token
	accessToken, err := getAccessToken()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting access token: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("Access token: %s...\n", accessToken[:20])

	// Connect to DuckDB
	db, err := sql.Open("duckdb", ":memory:")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening DuckDB: %v\n", err)
		os.Exit(1)
	}
	defer db.Close()

	// Load Azure/delta extensions
	fmt.Println("Installing and loading extensions...")
	_, err = db.Exec("INSTALL azure;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error installing azure: %v\n", err)
		os.Exit(1)
	}
	_, err = db.Exec("LOAD azure;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading azure: %v\n", err)
		os.Exit(1)
	}
	_, err = db.Exec("INSTALL delta;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error installing delta: %v\n", err)
		os.Exit(1)
	}
	_, err = db.Exec("LOAD delta;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading delta: %v\n", err)
		os.Exit(1)
	}

	// Create secret
	tablePath := fmt.Sprintf("abfss://%s@%s.dfs.core.windows.net/%s", container, account, relPath)
	secretSQL := fmt.Sprintf(`
		CREATE OR REPLACE SECRET azure_default (
			TYPE AZURE,
			PROVIDER access_token,
			ACCOUNT_NAME '%s',
			access_token '%s'
		);
	`, account, accessToken)

	fmt.Println("Creating Azure secret...")
	_, err = db.Exec(secretSQL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating secret: %v\n", err)
		os.Exit(1)
	}

	// Set configuration
	fmt.Println("Setting configuration...")
	_, err = db.Exec("SET threads=5;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error setting threads: %v\n", err)
		os.Exit(1)
	}
	_, err = db.Exec("SET azure_transport_option_type=curl;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error setting azure_transport_option_type: %v\n", err)
		os.Exit(1)
	}
	_, err = db.Exec("SET azure_read_transfer_concurrency=20;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error setting azure_read_transfer_concurrency: %v\n", err)
		os.Exit(1)
	}

	// Attach table
	tableAlias := "slowlogs_attached"
	fmt.Printf("Attaching table: %s as %s\n", tablePath, tableAlias)
	attachSQL := fmt.Sprintf("attach '%s' as %s(type delta);", tablePath, tableAlias)
	_, err = db.Exec(attachSQL)
	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "already exists") ||
			strings.Contains(strings.ToLower(err.Error()), "already attached") {
			fmt.Printf("Table %s already attached, continuing...\n", tableAlias)
		} else {
			fmt.Fprintf(os.Stderr, "Error attaching table: %v\n", err)
			os.Exit(1)
		}
	} else {
		fmt.Printf("Successfully attached table as %s\n", tableAlias)
	}

	// Query using attached table
	query := fmt.Sprintf(`
		SELECT id, conn_id, digest, time, query_time, mem_max, query, instance, succ,
		       (COALESCE(request_unit_read, 0) + COALESCE(request_unit_write, 0)) AS ru,
		       db, is_internal
		FROM %s
		WHERE time >= 1765867996.000000 AND time <= 1765869797.000000
		  AND date >= '20251216' AND date <= '20251216'
		  AND (is_internal = false OR is_internal IS NULL)
		ORDER BY time DESC
		LIMIT 100;
	`, tableAlias)

	fmt.Printf("\nQuery using attached table %s:\n", tableAlias)
	fmt.Println("Query:", query)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	rows, err := db.QueryContext(ctx, query)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing query: %v\n", err)
		os.Exit(1)
	}
	defer rows.Close()

	// Get column names
	columns, err := rows.Columns()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting columns: %v\n", err)
		os.Exit(1)
	}

	// Print header
	fmt.Println("\nResults:")
	for i, col := range columns {
		if i > 0 {
			fmt.Print(" | ")
		}
		fmt.Print(col)
	}
	fmt.Println()
	fmt.Println(strings.Repeat("-", 80))

	// Print rows
	rowCount := 0
	values := make([]interface{}, len(columns))
	valuePtrs := make([]interface{}, len(columns))
	for i := range values {
		valuePtrs[i] = &values[i]
	}

	for rows.Next() {
		err := rows.Scan(valuePtrs...)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error scanning row: %v\n", err)
			continue
		}

		for i, val := range values {
			if i > 0 {
				fmt.Print(" | ")
			}
			if val == nil {
				fmt.Print("NULL")
			} else {
				fmt.Print(val)
			}
		}
		fmt.Println()
		rowCount++
	}

	if err = rows.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "Error iterating rows: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\nTotal rows: %d\n", rowCount)
}