//go:build script

package main

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

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

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 expires in %d seconds\n", tokenResponse.ExpiresIn)
	return tokenResponse.AccessToken, nil
}

// DuckDBConnection wraps the DuckDB connection
type DuckDBConnection struct {
	DBX *sqlx.DB
	DB  *sql.DB
}

// DuckDBPool manages a pool of DuckDB connections using sync.Pool
type DuckDBPool struct {
	pool *sync.Pool
	once sync.Once
}

// NewDuckDBPool creates a new connection pool
func NewDuckDBPool() *DuckDBPool {
	return &DuckDBPool{
		pool: &sync.Pool{
			New: func() interface{} {
				conn := initializeDuckDB()
				if conn == nil {
					return nil
				}
				return conn
			},
		},
	}
}

// Get retrieves a connection from the pool
func (p *DuckDBPool) Get() *DuckDBConnection {
	fmt.Println("Getting connection from pool")
	conn := p.pool.Get()
	if conn == nil {
		return nil
	}
	return conn.(*DuckDBConnection)
}

// Put returns a connection to the pool
func (p *DuckDBPool) Put(conn *DuckDBConnection) {
	if conn != nil && conn.DB != nil {
		// Check if the connection is still valid
		if err := conn.DB.Ping(); err == nil {
			fmt.Println("Putting connection back to pool")
			p.pool.Put(conn)
		} else {
			// Connection is invalid, close it
			fmt.Println("Connection is invalid, closing it")
			conn.DB.Close()
			if conn.DBX != nil {
				conn.DBX.Close()
			}
		}
	}
}

// CloseAll closes all connections in the pool (for cleanup)
func (p *DuckDBPool) CloseAll() {
	// sync.Pool doesn't provide a way to iterate through all connections
	// This function is mainly for marking, actual cleanup is handled by GC
}

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

	// Create cache directory (same as duckdb_query.go)
	cacheDir := "cache/duckdb"
	if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
		err := os.MkdirAll(cacheDir, 0755)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error creating cache directory: %v\n", err)
			os.Exit(1)
		}
	}

	// Connect to DuckDB (same as duckdb_query.go)
	db, err := sql.Open("duckdb", ":memory:")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening DuckDB: %v\n", err)
		return nil
	}

	// Get DuckDB version (same as duckdb_query.go)
	var version string
	err = db.QueryRow("SELECT version()").Scan(&version)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting DuckDB version: %v\n", err)
		db.Close()
		return nil
	}
	fmt.Printf("DuckDB version: %s\n", version)

	// Set memory limit (same as duckdb_query.go)
	podMemLimitMBStr := os.Getenv("POD_MEMORY_LIMITS_MB")
	podMemLimitMB, parseErr := strconv.Atoi(podMemLimitMBStr)
	if parseErr != nil || podMemLimitMB <= 0 {
		fmt.Printf("Warning: Failed to parse POD_MEMORY_LIMITS_MB, using default 2048MB\n")
		podMemLimitMB = 2048
	}

	duckdbMemoryPercentStr := os.Getenv("DUCKDB_MEMORY_PERCENTAGE")
	duckdbMemoryPercent, parseErr := strconv.Atoi(duckdbMemoryPercentStr)
	if parseErr != nil || duckdbMemoryPercent <= 0 || duckdbMemoryPercent > 100 {
		fmt.Printf("Warning: Failed to parse DUCKDB_MEMORY_PERCENTAGE, using default 80%%\n")
		duckdbMemoryPercent = 80
	}

	duckdbMemoryLimitMB := int(float64(podMemLimitMB) * float64(duckdbMemoryPercent) / 100.0)
	if duckdbMemoryLimitMB < 2048 {
		duckdbMemoryLimitMB = 2048
		fmt.Printf("Calculated DuckDB memory limit too low, adjusting to minimum 2048MB\n")
	}

	duckdbMemoryLimit := fmt.Sprintf("%dMB", duckdbMemoryLimitMB)
	_, err = db.Exec(fmt.Sprintf("SET memory_limit='%s';", duckdbMemoryLimit))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error setting memory limit: %v\n", err)
		db.Close()
		return nil
	}
	fmt.Printf("Set DuckDB memory limit: %s\n", duckdbMemoryLimit)

	// Load extensions (same as duckdb_query.go)
	fmt.Println("Installing and loading extensions...")

	// Install and load cache_httpfs (same as duckdb_query.go)
	_, err = db.Exec(fmt.Sprintf("INSTALL cache_httpfs FROM community; LOAD cache_httpfs; SET cache_httpfs_cache_directory='%s';SET cache_httpfs_type='on_disk'; SET cache_httpfs_min_disk_bytes_for_cache=10000000000; ", cacheDir))
	if err != nil {
		fmt.Printf("Warning: Failed to install or load cache_httpfs extension, continuing without HTTP caching: %v\n", err)
		// Continue with initialization even if cache_httpfs fails
	} else {
		fmt.Println("Successfully loaded cache_httpfs extension")
	}

	// Install and load httpfs for S3 and GCS (same as duckdb_query.go)
	_, err = db.Exec("INSTALL httpfs; LOAD httpfs;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error installing/loading httpfs: %v\n", err)
		db.Close()
		return nil
	}

	// Install and load azure extension (same as duckdb_query.go)
	_, err = db.Exec("INSTALL azure; LOAD azure;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error installing/loading azure: %v\n", err)
		db.Close()
		return nil
	}

	// Set configuration (same as duckdb_query.go)
	fmt.Println("Setting configuration...")
	_, 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)
		db.Close()
		return nil
	}
	// Try lower concurrency to avoid SSL connection issues in concurrent queries
	// This is a workaround for curl backend SSL cert issues under concurrency
	azureConcurrency := os.Getenv("AZURE_READ_TRANSFER_CONCURRENCY")
	if azureConcurrency == "" {
		azureConcurrency = "5" // Lower default for concurrent queries
	}
	_, _ = db.Exec(fmt.Sprintf("set azure_read_transfer_concurrency=%s;", azureConcurrency))
	fmt.Printf("Set azure_read_transfer_concurrency=%s\n", azureConcurrency)

	threads := os.Getenv("THREADS")
	if threads == "" {
		threads = "5" // Lower default for concurrent queries
	}
	_, _ = db.Exec(fmt.Sprintf("set threads=%s;", threads))
	fmt.Printf("Set threads=%s\n", threads)

	// Install and load delta (same as duckdb_query.go)
	_, err = db.Exec("INSTALL delta; LOAD delta;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error installing/loading delta: %v\n", err)
		db.Close()
		return nil
	}

	// Set cache configuration (same as duckdb_query.go)
	_, err = db.Exec("SET enable_object_cache=true; SET enable_http_metadata_cache=true; SET enable_external_file_cache=true; SET parquet_metadata_cache=true;")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error setting cache configuration: %v\n", err)
		db.Close()
		return nil
	}

	// Create secret (same format as duckdb_query.go - without name)

	secretSQL := fmt.Sprintf(`
CREATE SECRET (
    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)
		db.Close()
		return nil
	}

	dbx := sqlx.NewDb(db, "duckdb")

	return &DuckDBConnection{
		DBX: dbx,
		DB:  db,
	}
}

func main() {
	var startTime, endTime string
	var hours int
	flag.StringVar(&startTime, "start-time", "1765867996.000000", "Start time (Unix timestamp)")
	flag.StringVar(&endTime, "end-time", "1765869797.000000", "End time (Unix timestamp)")
	flag.IntVar(&hours, "hours", 1, "Hours for metric/stats queries")
	flag.Parse()

	startTimeFloat, err := strconv.ParseFloat(startTime, 64)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing start-time: %v\n", err)
		os.Exit(1)
	}
	endTimeFloat, err := strconv.ParseFloat(endTime, 64)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error parsing end-time: %v\n", err)
		os.Exit(1)
	}
	beginHourTs := int64(startTimeFloat)

	fmt.Printf("Query parameters: start_time=%s, end_time=%s, hours=%d\n", startTime, endTime, hours)

	// Use connection pool to manage DuckDB connections
	// This allows connection reuse while keeping each query using an independent connection to avoid SSL concurrency issues
	tablePath := fmt.Sprintf("delta_scan('abfss://%s@%s.dfs.core.windows.net/%s')", container, account, relPath)

	// Create connection pool
	pool := NewDuckDBPool()

	// Clean up connection pool at program end (optional, sync.Pool manages automatically)
	defer func() {
		pool.CloseAll()
	}()

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

	// Run concurrent queries (same tableAlias, different queries)
	var wg sync.WaitGroup
	var mu sync.Mutex
	errors := make([]error, 0)

	// Add staggered delays to avoid simultaneous SSL connection initialization
	// This helps prevent curl backend SSL cert issues under concurrency
	queryStartDelay := os.Getenv("QUERY_START_DELAY_MS")
	delayMs := 0
	if queryStartDelay != "" {
		if parsed, err := strconv.Atoi(queryStartDelay); err == nil {
			delayMs = parsed
		}
	}

	// Option to serialize query execution with mutex (last resort for SSL issues)
	// Set SERIALIZE_QUERIES=1 to enable mutex protection
	serializeQueries := os.Getenv("SERIALIZE_QUERIES") == "1"
	var queryMu sync.Mutex // Mutex to serialize query execution if needed

	// 1. List query (like List function)
	wg.Add(1)
	go func() {
		defer wg.Done()
		// Get connection from pool

		if delayMs > 0 {
			time.Sleep(time.Duration(0*delayMs) * time.Millisecond)
		}
		conn := pool.Get()
		if conn == nil {
			mu.Lock()
			errors = append(errors, fmt.Errorf("List query failed: failed to get connection from pool"))
			mu.Unlock()
			return
		}
		// Return connection to pool after use
		defer pool.Put(conn)
		start := time.Now()
		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 >= %.6f AND time <= %.6f
			  AND (is_internal = false OR is_internal IS NULL)
			ORDER BY time DESC
			LIMIT 100;
		`, tablePath, startTimeFloat, endTimeFloat)

		type ResultRow struct {
			ID         sql.NullInt64   `db:"id"`
			ConnID     sql.NullString  `db:"conn_id"`
			Digest     sql.NullString  `db:"digest"`
			Time       sql.NullFloat64 `db:"time"`
			QueryTime  sql.NullFloat64 `db:"query_time"`
			MemMax     sql.NullFloat64 `db:"mem_max"`
			Query      sql.NullString  `db:"query"`
			Instance   sql.NullString  `db:"instance"`
			Succ       sql.NullBool    `db:"succ"`
			RU         sql.NullFloat64 `db:"ru"`
			DB         sql.NullString  `db:"db"`
			IsInternal sql.NullBool    `db:"is_internal"`
		}

		var results []ResultRow
		if serializeQueries {
			queryMu.Lock()
		}
		err := conn.DBX.SelectContext(ctx, &results, query)
		if serializeQueries {
			queryMu.Unlock()
		}
		elapsed := time.Since(start)
		mu.Lock()
		if err != nil {
			errors = append(errors, fmt.Errorf("List query failed: %w", err))
		} else {
			fmt.Printf("[List] Completed in %v, rows: %d\n", elapsed, len(results))
		}
		mu.Unlock()
	}()

	// 2. StatsMetric query (like StatsMetric function)
	wg.Add(1)
	go func() {
		defer wg.Done()

		if delayMs > 0 {
			time.Sleep(time.Duration(1*delayMs) * time.Millisecond)
		}
		// Get connection from pool
		conn := pool.Get()
		if conn == nil {
			mu.Lock()
			errors = append(errors, fmt.Errorf("StatsMetric query failed: failed to get connection from pool"))
			mu.Unlock()
			return
		}
		// Return connection to pool after use
		defer pool.Put(conn)
		start := time.Now()
		// Calculate bucket size (same as StatsMetric)
		bucketSeconds := int64(hours * 3600 / 120)
		if bucketSeconds < 60 {
			bucketSeconds = 60
		}
		query := fmt.Sprintf(`
			SELECT
				FLOOR(time / %d) * %d AS bucket_ts,
				COUNT(*) as count
			FROM %s
			WHERE time >= %d AND time <= %d AND (is_internal = false OR is_internal IS NULL)
			GROUP BY bucket_ts ORDER BY bucket_ts ASC
		`, bucketSeconds, bucketSeconds, tablePath, beginHourTs, beginHourTs+int64(hours)*3600)

		type StatsMetricResult struct {
			BucketTS float64 `db:"bucket_ts"`
			Count    int64   `db:"count"`
		}
		var results []StatsMetricResult
		if serializeQueries {
			queryMu.Lock()
		}
		err := conn.DBX.SelectContext(ctx, &results, query)
		if serializeQueries {
			queryMu.Unlock()
		}
		elapsed := time.Since(start)
		mu.Lock()
		if err != nil {
			errors = append(errors, fmt.Errorf("StatsMetric query failed: %w", err))
		} else {
			fmt.Printf("[StatsMetric] Completed in %v, buckets: %d\n", elapsed, len(results))
		}
		mu.Unlock()
	}()

	// 3. Stats query (like Stats function)
	wg.Add(1)
	go func() {
		defer wg.Done()
		// Get connection from pool

		if delayMs > 0 {
			time.Sleep(time.Duration(2*delayMs) * time.Millisecond)
		}
		conn := pool.Get()
		if conn == nil {
			mu.Lock()
			errors = append(errors, fmt.Errorf("Stats query failed: failed to get connection from pool"))
			mu.Unlock()
			return
		}
		// Return connection to pool after use
		defer pool.Put(conn)
		start := time.Now()
		query := fmt.Sprintf(`
			SELECT
				avg(disk_max) as avg_disk,
				avg(query_time) as avg_latency,
				avg(mem_max) as avg_memory,
				count(*) as count,
				first(db) as db,
				max(disk_max) as max_disk,
				max(query_time) as max_latency,
				max(mem_max) as max_memory,
				digest as sql_digest,
				first(query) as sql_text,
				sum(query_time) as sum_latency,
				sum(mem_max) as sum_memory,
				sum(disk_max) as sum_disk,
				CASE
				    WHEN first(query) ILIKE 'select%%' THEN 'Select'
				    WHEN first(query) ILIKE 'insert%%' THEN 'Insert'
				    WHEN first(query) ILIKE 'update%%' THEN 'Update'
				    WHEN first(query) ILIKE 'delete%%' THEN 'Delete'
				    ELSE 'Others'
				END as statement_type
			FROM %s
			WHERE time >= %d AND time <= %d AND (is_internal = false OR is_internal IS NULL)
			GROUP BY digest
			ORDER BY sum_latency DESC
			LIMIT 10
		`, tablePath, beginHourTs, beginHourTs+int64(hours)*3600)

		type StatsResult struct {
			AvgDisk       sql.NullFloat64 `db:"avg_disk"`
			AvgLatency    sql.NullFloat64 `db:"avg_latency"`
			AvgMemory     sql.NullFloat64 `db:"avg_memory"`
			Count         int             `db:"count"`
			DB            sql.NullString  `db:"db"`
			MaxDisk       sql.NullFloat64 `db:"max_disk"`
			MaxLatency    sql.NullFloat64 `db:"max_latency"`
			MaxMemory     sql.NullFloat64 `db:"max_memory"`
			SQLDigest     sql.NullString  `db:"sql_digest"`
			SQLText       sql.NullString  `db:"sql_text"`
			SumLatency    sql.NullFloat64 `db:"sum_latency"`
			SumMemory     sql.NullFloat64 `db:"sum_memory"`
			SumDisk       sql.NullFloat64 `db:"sum_disk"`
			StatementType sql.NullString  `db:"statement_type"`
		}
		var results []StatsResult
		if serializeQueries {
			queryMu.Lock()
		}
		err := conn.DBX.SelectContext(ctx, &results, query)
		if serializeQueries {
			queryMu.Unlock()
		}
		elapsed := time.Since(start)
		mu.Lock()
		if err != nil {
			errors = append(errors, fmt.Errorf("Stats query failed: %w", err))
		} else {
			fmt.Printf("[Stats] Completed in %v, rows: %d\n", elapsed, len(results))
		}
		mu.Unlock()
	}()

	// 4. Databases query (like Databases function)
	wg.Add(1)
	go func() {
		defer wg.Done()
		// Get connection from pool

		if delayMs > 0 {
			time.Sleep(time.Duration(3*delayMs) * time.Millisecond)
		}

		conn := pool.Get()
		if conn == nil {
			mu.Lock()
			errors = append(errors, fmt.Errorf("Databases query failed: failed to get connection from pool"))
			mu.Unlock()
			return
		}
		// Return connection to pool after use
		defer pool.Put(conn)
		start := time.Now()
		beginDate := time.Unix(beginHourTs, 0).UTC().Format("20060102")
		endDate := time.Unix(beginHourTs+int64(hours)*3600, 0).UTC().Format("20060102")
		query := fmt.Sprintf(`
			SELECT DISTINCT db
			FROM %s
			WHERE time >= %.6f AND time <= %.6f
			  AND date >= '%s' AND date <= '%s'
			  AND db IS NOT NULL
			  AND db != ''
			  AND (is_internal = false OR is_internal IS NULL)
			ORDER BY db
		`, tablePath, startTimeFloat, endTimeFloat, beginDate, endDate)

		type DatabaseResult struct {
			DB sql.NullString `db:"db"`
		}
		var results []DatabaseResult
		if serializeQueries {
			queryMu.Lock()
		}
		err := conn.DBX.SelectContext(ctx, &results, query)
		if serializeQueries {
			queryMu.Unlock()
		}
		elapsed := time.Since(start)
		mu.Lock()
		if err != nil {
			errors = append(errors, fmt.Errorf("Databases query failed: %w", err))
		} else {
			databases := make([]string, 0, len(results))
			for _, r := range results {
				if r.DB.Valid {
					databases = append(databases, r.DB.String)
				}
			}
			fmt.Printf("[Databases] Completed in %v, databases: %d %v\n", elapsed, len(databases), databases)
		}
		mu.Unlock()
	}()

	// Wait for all queries to complete
	fmt.Printf("\nStarting concurrent queries using connection pool...")
	if delayMs > 0 {
		fmt.Printf(" (with %dms staggered delays)", delayMs)
	}
	if serializeQueries {
		fmt.Printf(" (serialized with mutex)")
	}
	fmt.Println()
	wg.Wait()

	// Report results
	fmt.Println("\n=== Concurrent Query Results ===")
	if len(errors) > 0 {
		fmt.Println("Errors occurred:")
		for _, err := range errors {
			fmt.Printf("  - %v\n", err)
		}
		os.Exit(1)
	} else {
		fmt.Println("All queries completed successfully!")
	}
}