package duckdb

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

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	_ "github.com/duckdb/duckdb-go/v2"
	"github.com/jmoiron/sqlx"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type QueryDuckDBFetcherInterface interface {
	// SelectSlowlogInDL executes a query and scans results into the provided destination slice.
	// The dest parameter should be a pointer to a slice of structs with appropriate db tags.
	SelectSlowlogInDL(ctx context.Context, dest interface{}, fullQuery string) error

	// SelectSlowlogInDLWithoutFilterPushdown executes a query with a specific table alias and scans results into dest.
	// The dest parameter should be a pointer to a slice of structs with appropriate db tags.
	SelectSlowlogInDLWithoutFilterPushdown(ctx context.Context, dest interface{}, tableAlias, table, fullQuery string) error

	// Close closes the database connection.
	Close(ctx context.Context) error
}

type QueryDuckDBFetcher struct {
	vendor         string
	region         string
	db             *sql.DB
	dbx            *sqlx.DB
	mu             sync.RWMutex
	attachedTables map[string]string
	pool           *DuckDBPool // connection pool
	usePool        bool        // whether to use connection pool
}

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

	// lastUsed is updated by the pool on Put/Get for idle timeout eviction.
	lastUsed time.Time
}

// DuckDBPool manages DuckDB connections with an idle timeout.
type DuckDBPool struct {
	vendor    string
	region    string
	maxIdle   time.Duration
	mu        sync.Mutex
	idleConns []*DuckDBConnection

	closeCh   chan struct{}
	closeOnce sync.Once
}

var sqlOpen = sql.Open

// Global connection pool manager (singleton pattern)
var (
	poolManager   = make(map[string]*DuckDBPool)
	poolManagerMu sync.RWMutex
)

// getPoolKey generates a unique key for vendor+region combination
func getPoolKey(vendor, region string) string {
	return fmt.Sprintf("%s:%s", vendor, region)
}

// GetOrCreateDuckDBPool returns a singleton DuckDBPool for the given vendor and region
func GetOrCreateDuckDBPool(vendor, region string) *DuckDBPool {
	key := getPoolKey(vendor, region)

	// Try to get existing pool (read lock)
	poolManagerMu.RLock()
	if pool, exists := poolManager[key]; exists {
		poolManagerMu.RUnlock()
		log.Info("Reusing existing DuckDB connection pool", zap.String("vendor", vendor), zap.String("region", region))
		return pool
	}
	poolManagerMu.RUnlock()

	// Create new pool (write lock)
	poolManagerMu.Lock()
	defer poolManagerMu.Unlock()

	// Double-check after acquiring write lock
	if pool, exists := poolManager[key]; exists {
		log.Info("Reusing existing DuckDB connection pool (after lock)", zap.String("vendor", vendor), zap.String("region", region))
		return pool
	}

	// Create new pool
	pool := NewDuckDBPool(vendor, region)
	poolManager[key] = pool
	log.Info("Created new DuckDB connection pool (singleton)", zap.String("vendor", vendor), zap.String("region", region))
	return pool
}

func NewQueryDuckDBFetcher(ctx context.Context, vendor, region string) (*QueryDuckDBFetcher, error) {
	s := &QueryDuckDBFetcher{
		vendor:         vendor,
		region:         region,
		attachedTables: make(map[string]string),
		usePool:        false, // do not use connection pool by default, maintain backward compatibility
	}

	err := s.Init(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize SlowQueryDuckDBFetcher: %w", err)
	}
	return s, err
}

// NewQueryDuckDBFetcherWithPool creates a QueryDuckDBFetcher that uses connection pool (singleton)
func NewQueryDuckDBFetcherWithPool(ctx context.Context, vendor, region string) (*QueryDuckDBFetcher, error) {
	s := &QueryDuckDBFetcher{
		vendor:         vendor,
		region:         region,
		attachedTables: make(map[string]string),
		usePool:        true,
	}

	// Get or create singleton connection pool
	pool := GetOrCreateDuckDBPool(vendor, region)
	s.pool = pool

	log.Info("QueryDuckDBFetcher initialized with connection pool (singleton)", zap.String("vendor", vendor), zap.String("region", region))
	return s, nil
}

// NewDuckDBPool creates a new connection pool
func NewDuckDBPool(vendor, region string) *DuckDBPool {
	pool := &DuckDBPool{
		vendor: vendor,
		region: region,
	}
	pool.maxIdle = 5 * time.Minute
	pool.closeCh = make(chan struct{})

	// Background goroutine to evict idle connections deterministically.
	pool.startJanitor()
	return pool
}

// Get retrieves a connection from the pool
func (p *DuckDBPool) Get() *DuckDBConnection {
	log.Info("Getting connection from pool")

	p.mu.Lock()
	now := time.Now()

	// Evict expired idle conns eagerly so we don't rely on Go GC.
	// nolint:prealloc
	var kept, expired []*DuckDBConnection
	for _, conn := range p.idleConns {
		if conn == nil || conn.DB == nil {
			continue
		}
		if now.Sub(conn.lastUsed) >= p.maxIdle {
			expired = append(expired, conn)
			continue
		}
		kept = append(kept, conn)
	}
	p.idleConns = kept

	if len(p.idleConns) > 0 {
		conn := p.idleConns[len(p.idleConns)-1]
		p.idleConns = p.idleConns[:len(p.idleConns)-1]
		p.mu.Unlock()

		if conn == nil {
			log.Warn("Got nil connection from idle list")
			return nil
		}
		conn.lastUsed = now

		if len(expired) > 0 {
			for _, c := range expired {
				closeDuckDBConnection(c)
			}
		}

		log.Info("Got connection from pool")
		return conn
	}
	p.mu.Unlock()

	// Close expired outside the lock to avoid blocking queries.
	for _, c := range expired {
		closeDuckDBConnection(c)
	}

	log.Info("No idle connection available, creating new one")
	conn := p.initializeConnection()
	if conn != nil {
		conn.lastUsed = now
		log.Info("Created new connection for pool")
	}
	return conn
}

// Put returns a connection to the pool
func (p *DuckDBPool) Put(conn *DuckDBConnection) {
	if conn == nil {
		log.Warn("Attempted to put nil connection to pool")
		return
	}
	if conn.DB == nil {
		log.Warn("Attempted to put connection with nil DB to pool")
		return
	}
	log.Info("Putting connection back to pool")

	conn.lastUsed = time.Now()
	p.mu.Lock()
	p.idleConns = append(p.idleConns, conn)
	p.mu.Unlock()
}

func closeDuckDBConnection(conn *DuckDBConnection) {
	if conn == nil {
		return
	}
	// Ensure Close doesn't panic on nil sub-components.
	if conn.DBX != nil {
		_ = conn.DBX.Close()
	}
	if conn.DB != nil {
		_ = conn.DB.Close()
	}
	conn.DBX = nil
	conn.DB = nil
}

func (p *DuckDBPool) startJanitor() {
	// Prevent multiple janitors for the same pool instance.
	go func() {
		ticker := time.NewTicker(time.Minute)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				p.evictExpired()
			case <-p.closeCh:
				return
			}
		}
	}()
}

func (p *DuckDBPool) evictExpired() {
	now := time.Now()

	p.mu.Lock()
	//nolint:prealloc
	var expired, kept []*DuckDBConnection
	for _, conn := range p.idleConns {
		if conn == nil || conn.DB == nil {
			continue
		}
		if now.Sub(conn.lastUsed) >= p.maxIdle {
			expired = append(expired, conn)
			continue
		}
		kept = append(kept, conn)
	}
	p.idleConns = kept
	p.mu.Unlock()

	for _, c := range expired {
		log.Info("Closing idle DuckDB connection",
			zap.Duration("idle", now.Sub(c.lastUsed)),
			zap.String("vendor", p.vendor),
			zap.String("region", p.region),
		)
		closeDuckDBConnection(c)
	}
}

// CloseAll closes the pool janitor and releases all idle connections.
func (p *DuckDBPool) CloseAll() {
	p.closeOnce.Do(func() {
		close(p.closeCh)
	})

	p.mu.Lock()
	idle := p.idleConns
	p.idleConns = nil
	p.mu.Unlock()

	for _, c := range idle {
		closeDuckDBConnection(c)
	}
}

// Secret templates for different cloud providers
var (
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	awsSecret = `
CREATE SECRET (
    TYPE s3,
    PROVIDER credential_chain,
    ENDPOINT 's3.%s.amazonaws.com',
    REFRESH auto,
    REGION '%s'
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	awsReplaceSecret = `
CREATE OR REPLACE SECRET (
    TYPE s3,
    PROVIDER credential_chain,
    ENDPOINT 's3.%s.amazonaws.com',
    REFRESH auto,
    REGION '%s'
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	aliSecret = `
CREATE SECRET (
    TYPE s3,
    PROVIDER credential_chain,
    ENDPOINT 'oss-%s-internal.aliyuncs.com',
    REFRESH auto,
    REGION '%s'
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	aliReplaceSecret = `
CREATE OR REPLACE SECRET (
    TYPE s3,
    PROVIDER credential_chain,
    ENDPOINT 'oss-%s-internal.aliyuncs.com',
    REFRESH auto,
    REGION '%s'
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	gcsSecret = `
CREATE SECRET (
    TYPE gcs,
    PROVIDER credential_chain
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	gcsReplaceSecret = `
CREATE OR REPLACE SECRET (
    TYPE gcs,
    PROVIDER credential_chain
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	azureSecret = `
CREATE SECRET (
    TYPE azure,
    PROVIDER access_token,
	ACCOUNT_NAME '%s',
	access_token '%s'
);
`
	// nolint:gosec // G101: This is a DDL statement to create a secret *configuration*, not a hardcoded credential.
	azureReplaceSecret = `
CREATE OR REPLACE SECRET (
    TYPE azure,
    PROVIDER access_token,
	ACCOUNT_NAME '%s',
	access_token '%s'
);
`
)

const cacheDir = "cache/duckdb"
const duckDBAssumeRoleSessionName = "nextgen-reporter-duckdb"

func getAWSAccessToken(region, roleArn string) (*aws.Credentials, error) {
	ctx := context.Background()

	// Load AWS config using default credential chain
	cfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(region))
	if err != nil {
		log.Error("failed to load AWS config", zap.Error(err))
		return nil, fmt.Errorf("failed to load AWS config: %w", err)
	}

	if roleArn = strings.TrimSpace(roleArn); roleArn != "" {
		// Get current IAM identity to check if we're already using the target role
		stsClient := sts.NewFromConfig(cfg)
		identity, err := stsClient.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
		if err != nil {
			log.Warn("failed to get AWS caller identity, will try to assume role", zap.Error(err))
		} else {
			// If we're already using the target role, just return current credentials
			log.Info("get AWS caller identity", zap.String("role_arn", roleArn), zap.String("identity", *identity.Arn))
			if identity.Arn != nil && strings.Contains(*identity.Arn, roleArn[strings.LastIndex(roleArn, "/")+1:]) {
				log.Info("Skipping role assumption as current role matches target role",
					zap.String("role_arn", roleArn))
				creds, err := cfg.Credentials.Retrieve(ctx)
				if err != nil {
					return nil, fmt.Errorf("failed to retrieve current credentials: %w", err)
				}
				return &creds, nil
			}
		}

		// If we're not already using the target role, proceed with role assumption
		provider := stscreds.NewAssumeRoleProvider(stsClient, roleArn, func(o *stscreds.AssumeRoleOptions) {
			o.RoleSessionName = duckDBAssumeRoleSessionName
		})
		credCache := aws.NewCredentialsCache(provider)
		credentials, err := credCache.Retrieve(ctx)
		if err != nil {
			log.Error("failed to assume AWS role for DuckDB",
				zap.String("role_arn", roleArn),
				zap.Error(err))
			return nil, fmt.Errorf("failed to retrieve assumed credentials: %w", err)
		}

		maskedAccessKey := credentials.AccessKeyID
		if len(maskedAccessKey) > 8 {
			maskedAccessKey = maskedAccessKey[:8] + "..."
		}
		log.Info("Successfully assumed AWS role for DuckDB",
			zap.String("role_arn", roleArn),
			zap.String("access_key_id", maskedAccessKey),
			zap.Bool("has_session_token", credentials.SessionToken != ""),
			zap.Time("expires", credentials.Expires))

		return &credentials, nil
	}

	// Retrieve credentials from the credential chain
	credentials, err := cfg.Credentials.Retrieve(ctx)
	if err != nil {
		log.Error("failed to retrieve AWS credentials from credential chain", zap.Error(err))
		return nil, fmt.Errorf("failed to retrieve AWS credentials: %w", err)
	}

	maskedAccessKey := credentials.AccessKeyID
	if len(maskedAccessKey) > 8 {
		maskedAccessKey = maskedAccessKey[:8] + "..."
	}
	log.Info("Successfully retrieved AWS credentials from SDK",
		zap.String("access_key_id", maskedAccessKey),
		zap.Bool("has_session_token", credentials.SessionToken != ""),
		zap.Time("expires", credentials.Expires))

	return &credentials, nil

}

// getAzureAccessToken retrieves Azure access token using Workload Identity
func getAzureAccessToken() (string, error) {

	federatedTokenFile := os.Getenv("AZURE_FEDERATED_TOKEN_FILE")
	if federatedTokenFile == "" {
		return "", fmt.Errorf("env 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("env 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 returns 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("response does not contain access token")
	}

	log.Info("Successfully obtained Azure access token", zap.String("client_id", clientID), zap.String("tenant_id", tenantID))
	return tokenResponse.AccessToken, nil
}

// initializeConnection initializes a single DuckDB connection (for connection pool use)
func (p *DuckDBPool) initializeConnection() *DuckDBConnection {
	return initializeDuckDBConnection(p.vendor, p.region)
}

// initializeDuckDBConnection initializes a single DuckDB connection
//
//nolint:gocyclo // Complex initialization function with multiple vendor-specific configurations
func initializeDuckDBConnection(vendor, region string) *DuckDBConnection {
	if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
		err := os.MkdirAll(cacheDir, 0755)
		if err != nil {
			log.Error("Failed to create cache directory", zap.String("cacheDir", cacheDir), zap.Error(err))
			return nil
		}
	}

	newDB, err := sqlOpen("duckdb", ":memory:")
	if err != nil {
		log.Error("Failed to open duckdb", zap.Error(err))
		return nil
	}

	podMemLimitMBStr := os.Getenv("POD_MEMORY_LIMITS_MB")
	podMemLimitMB, parseErr := strconv.Atoi(podMemLimitMBStr)
	if parseErr != nil || podMemLimitMB <= 0 {
		log.Warn("Failed to parse POD_MEMORY_LIMITS_MB env var", zap.Error(parseErr), zap.String("POD_MEMORY_LIMITS_MB", podMemLimitMBStr))
		podMemLimitMB = 2048
	}

	duckdbMemoryPercentStr := os.Getenv("DUCKDB_MEMORY_PERCENTAGE")
	duckdbMemoryPercent, parseErr := strconv.Atoi(duckdbMemoryPercentStr)
	if parseErr != nil || duckdbMemoryPercent <= 0 || duckdbMemoryPercent > 100 {
		log.Warn("Failed to parse DUCKDB_MEMORY_PERCENTAGE env var or invalid value, using default 80%", zap.Error(parseErr), zap.String("DUCKDB_MEMORY_PERCENTAGE", duckdbMemoryPercentStr))
		duckdbMemoryPercent = 80
	}

	duckdbMemoryLimitMB := int(float64(podMemLimitMB) * float64(duckdbMemoryPercent) / 100.0)

	if duckdbMemoryLimitMB < 2048 {
		duckdbMemoryLimitMB = 2048
		log.Info("Calculated DuckDB memory limit too low, adjusting to minimum", zap.Int("calculatedMB", duckdbMemoryLimitMB))
	}

	duckdbMemoryLimit := fmt.Sprintf("%dMB", duckdbMemoryLimitMB)

	_, err = newDB.Exec(fmt.Sprintf("SET memory_limit='%s';", duckdbMemoryLimit))
	if err != nil {
		newDB.Close()
		log.Error("Failed to set DuckDB memory_limit", zap.String("limit", duckdbMemoryLimit), zap.Error(err))
		return nil
	}
	log.Info("Set DuckDB memory limit", zap.String("limit", duckdbMemoryLimit),
		zap.Int("pod_memory_limit_mb", podMemLimitMB),
		zap.Int("duckdbMemoryPercent", duckdbMemoryPercent))

	_, err = newDB.Exec(fmt.Sprintf("INSTALL cache_httpfs FROM community; LOAD cache_httpfs; SET cache_httpfs_cache_directory='%s';SET cache_httpfs_type='on_disk'; ", cacheDir))
	if err != nil {
		log.Warn("Failed to install or load cache_httpfs extension, continuing without HTTP caching",
			zap.Error(err),
			zap.String("cacheDir", cacheDir))
		// Continue with initialization even if cache_httpfs fails
		// The application can still function without HTTP caching
	} else {
		log.Info("Successfully loaded cache_httpfs extension", zap.String("cacheDir", cacheDir))
	}
	// Install and load httpfs for S3 and GCS
	_, err = newDB.Exec("INSTALL httpfs; LOAD httpfs;")
	if err != nil {
		newDB.Close()
		log.Error("Failed to install or load httpfs", zap.Error(err))
		return nil
	}

	// Install and load azure extension for Azure Blob Storage
	_, err = newDB.Exec("INSTALL azure; LOAD azure;")
	if err != nil {
		newDB.Close()
		log.Error("Failed to install or load azure extension", zap.Error(err))
		return nil
	}

	_, err = newDB.Exec("set azure_transport_option_type=curl;")
	if err != nil {
		newDB.Close()
		log.Error("Failed to set azure_transport_option_type", zap.Error(err))
		return nil
	}
	azureConcurrency := os.Getenv("AZURE_READ_TRANSFER_CONCURRENCY")
	if azureConcurrency == "" {
		azureConcurrency = "5"
	}
	_, _ = newDB.Exec(fmt.Sprintf("set azure_read_transfer_concurrency=%s;", azureConcurrency))

	threads := os.Getenv("THREADS")
	if threads == "" {
		threads = "5"
	}
	_, _ = newDB.Exec(fmt.Sprintf("set threads=%s;", threads))

	_, err = newDB.Exec("INSTALL delta; LOAD delta;")
	if err != nil {
		newDB.Close()
		log.Error("Failed to install or load delta", zap.Error(err))
		return nil
	}

	_, err = newDB.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 {
		newDB.Close()
		log.Error("Failed to set cache configuration", zap.Error(err))
		return nil
	}

	// Configure secrets based on vendor
	switch strings.ToLower(vendor) {
	case "aws":
		_, err = newDB.Exec(fmt.Sprintf("SET s3_region='%s';", region))
		if err != nil {
			newDB.Close()
			log.Error("Failed to set s3_region", zap.Error(err))
			return nil
		}
		roleArn := os.Getenv("AWS_ROLE_ARN")
		credentials, err := getAWSAccessToken(region, roleArn)
		if err != nil {
			newDB.Close()
			log.Error("Failed to get AWS access token", zap.Error(err))
			return nil
		}
		var awsSecretSQL string
		if credentials.SessionToken != "" {
			awsSecretSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					SESSION_TOKEN '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, credentials.SessionToken, region)
		} else {
			awsSecretSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, region)
		}
		_, err = newDB.Exec(awsSecretSQL)
		if err != nil {
			newDB.Close()
			log.Error("Failed to set AWS secret", zap.Error(err))
			return nil
		}
	case "alicloud":
		_, err = newDB.Exec(fmt.Sprintf("SET s3_region='%s';", region))
		if err != nil {
			newDB.Close()
			log.Error("Failed to set s3_region", zap.Error(err))
			return nil
		}
		_, err = newDB.Exec("SET s3_url_style='vhost';")
		if err != nil {
			newDB.Close()
			log.Error("Failed to set s3_url_style", zap.Error(err))
			return nil
		}
		_, err = newDB.Exec(fmt.Sprintf(aliSecret, region, region))
		if err != nil {
			newDB.Close()
			log.Error("Failed to set OSS secret", zap.Error(err))
			return nil
		}
	case "gcp":
		// GCS does not use 'region' in the same way as S3 for endpoint,
		// the credential_chain should handle it.
		// HMAC keys are required for GCS with S3-like API.
		// DuckDB's GCS secret type still uses the httpfs extension.
		_, err = newDB.Exec(gcsSecret)
		if err != nil {
			newDB.Close()
			log.Error("Failed to set GCP secret", zap.Error(err))
			return nil
		}
	case "azure":
		// Azure blob storage uses access_code with Workload Identity token
		// Get access token and configure DuckDB secret
		accessToken, err := getAzureAccessToken()
		if err != nil {
			newDB.Close()
			log.Error("Failed to get Azure access token", zap.Error(err))
			return nil
		}
		_, err = newDB.Exec(fmt.Sprintf(azureSecret, os.Getenv("AZURE_STORAGE_ACCOUNT"), accessToken))
		if err != nil {
			newDB.Close()
			log.Error("Failed to set Azure secret", zap.Error(err))
			return nil
		}
	case "local":
		// Do nothing
	default:
		newDB.Close()
		log.Error("Unsupported vendor", zap.String("vendor", vendor))
		return nil
	}

	newDBx := sqlx.NewDb(newDB, "duckdb")

	return &DuckDBConnection{
		DBX: newDBx,
		DB:  newDB,
	}
}

func (s *QueryDuckDBFetcher) Init(ctx context.Context) error {
	conn := initializeDuckDBConnection(s.vendor, s.region)
	if conn == nil {
		return fmt.Errorf("failed to initialize DuckDB connection")
	}

	s.db = conn.DB
	s.dbx = conn.DBX
	s.attachedTables = map[string]string{}

	log.Info("SlowQueryDuckDBFetcher initialized")
	return nil
}

// ExecDDL runs DDL statements (for example CREATE VIEW in integration tests).
func (s *QueryDuckDBFetcher) ExecDDL(ctx context.Context, query string) error {
	if s.usePool && s.pool != nil {
		conn := s.pool.Get()
		if conn == nil {
			return fmt.Errorf("failed to get connection from pool")
		}
		defer s.pool.Put(conn)
		_, err := conn.DBX.ExecContext(ctx, query)
		return err
	}
	_, err := s.dbx.ExecContext(ctx, query)
	return err
}

func (s *QueryDuckDBFetcher) Close(ctx context.Context) error {
	// For connection pool mode, we don't close the shared pool here
	// as it's a singleton shared across multiple QueryDuckDBFetcher instances
	// The pool will be closed when the application shuts down
	if s.usePool && s.pool != nil {
		log.Info("QueryDuckDBFetcher closed (pool remains open for other instances)", zap.String("vendor", s.vendor), zap.String("region", s.region))
		// Don't close the shared pool - it's a singleton
		// s.pool.CloseAll() // Commented out to preserve singleton pool
	}

	if s.dbx != nil {
		err := s.dbx.Close()
		if err != nil {
			return err
		}
	}
	if s.db != nil {
		err := s.db.Close()
		if err != nil {
			return err
		}
	}

	runtime.GC()
	debug.FreeOSMemory()
	return nil
}

func (s *QueryDuckDBFetcher) SelectSlowlogInDL(ctx context.Context, dest interface{}, fullQuery string) error {
	start := time.Now()
	defer func() {
		end := time.Now()
		log.Info("SlowQueryDuckDBFetcher SelectSlowlogInDL called",
			zap.String("fullQuery", fullQuery),
			zap.Int("cost_ms", int(end.Sub(start).Milliseconds())),
		)
	}()

	if s.usePool && s.pool != nil {
		// Using connection pool
		conn := s.pool.Get()
		if conn == nil {
			return fmt.Errorf("failed to get connection from pool")
		}
		defer s.pool.Put(conn)

		err := conn.DBX.SelectContext(ctx, dest, fullQuery)
		if err != nil {
			return fmt.Errorf("failed to query and scan with sqlx: %w", err)
		}
		return nil
	}

	// Using single connection mode
	err := s.dbx.SelectContext(ctx, dest, fullQuery)
	if err != nil {
		return fmt.Errorf("failed to query and scan with sqlx: %w", err)
	}

	return nil
}

func (s *QueryDuckDBFetcher) SelectSlowlogInDLWithoutFilterPushdown(ctx context.Context, dest interface{}, tableAlias, table, fullQuery string) error {
	start := time.Now()
	defer func() {
		end := time.Now()
		cost_ms := int(end.Sub(start).Milliseconds())
		if cost_ms > 15000 {
			log.Warn("SlowQueryDuckDBFetcher SelectSlowlogInDLWithoutFilterPushdown called",
				zap.String("fullQuery", fullQuery),
				zap.Int("cost_ms", cost_ms),
			)
		} else {
			log.Info("SlowQueryDuckDBFetcher SelectSlowlogInDLWithoutFilterPushdown called",
				zap.String("fullQuery", fullQuery),
				zap.Int("cost_ms", cost_ms),
			)
		}
	}()

	// Handle map slice destination
	destValue := reflect.ValueOf(dest)
	if destValue.Kind() != reflect.Ptr {
		return fmt.Errorf("dest must be a pointer")
	}

	destElem := destValue.Elem()
	if destElem.Kind() != reflect.Slice {
		return fmt.Errorf("dest must be a pointer to slice")
	}

	elemType := destElem.Type().Elem()
	// There is a bug in duckdb that it cannot handle []map[string]interface{}, so we need to handle it manually
	if elemType.String() == "map[string]interface {}" {
		log.Info("Using custom map slice handler for []map[string]interface{}")
		return s.selectIntoMapSliceWithPool(ctx, dest, tableAlias, table, fullQuery)
	}

	// use pooled connection mode
	if s.usePool && s.pool != nil {
		conn := s.pool.Get()
		if conn == nil {
			return fmt.Errorf("failed to get connection from pool")
		}
		defer s.pool.Put(conn)

		// attach table on connection
		err := s.attachTableOnConnection(conn, tableAlias, table)
		if err != nil {
			return err
		}

		// For struct slices, use normal sqlx behavior
		err = conn.DBX.SelectContext(ctx, dest, fullQuery)
		if err != nil {
			if strings.Contains(err.Error(), "The provided token has expired.") ||
				strings.Contains(err.Error(), "AuthenticationFailed") ||
				strings.Contains(err.Error(), "Credential should be updated") {
				log.Warn("Credential error detected, attempting to refresh credential chain", zap.Error(err))
				err = s.refreshCredentialChainOnConnection(ctx, conn)
				if err != nil {
					return fmt.Errorf("failed to refresh credential chain after query failed: %w", err)
				}

				log.Info("Credential chain refreshed successfully, retrying query")
				err = conn.DBX.SelectContext(ctx, dest, fullQuery)
				if err != nil {
					return fmt.Errorf("failed to query and scan with sqlx after credential refresh: %w", err)
				}
				return nil
			} else if strings.Contains(err.Error(), "No files in log segment") {
				log.Warn("No files in log segment", zap.Error(err))
				return nil
			} else if strings.Contains(err.Error(), "context canceled") {
				log.Warn("Context canceled", zap.Error(err))
				return nil
			}

			return fmt.Errorf("failed to query and scan with sqlx: %w", err)
		}
		return nil
	}

	// use single connection mode
	err := s.attachTable(tableAlias, table)
	if err != nil {
		return err
	}

	// For struct slices, use normal sqlx behavior
	err = s.dbx.SelectContext(ctx, dest, fullQuery)
	if err != nil {
		if strings.Contains(err.Error(), "The provided token has expired.") ||
			strings.Contains(err.Error(), "AuthenticationFailed") ||
			strings.Contains(err.Error(), "Credential should be updated") {
			log.Warn("Credential error detected, attempting to refresh credential chain", zap.Error(err))
			err = s.refreshCredentialChain(ctx)
			if err != nil {
				return fmt.Errorf("failed to refresh credential chain after query failed: %w", err)
			}

			log.Info("Credential chain refreshed successfully, retrying query")
			err = s.dbx.SelectContext(ctx, dest, fullQuery)
			if err != nil {
				return fmt.Errorf("failed to query and scan with sqlx after credential refresh: %w", err)
			}
			return nil
		} else if strings.Contains(err.Error(), "No files in log segment") {
			log.Warn("No files in log segment", zap.Error(err))
			return nil
		}
		return fmt.Errorf("failed to query and scan with sqlx: %w", err)
	}

	return nil
}

// selectIntoMapSliceWithPool handles scanning into []map[string]interface{} with connection pool
func (s *QueryDuckDBFetcher) selectIntoMapSliceWithPool(ctx context.Context, dest interface{}, tableAlias, table, fullQuery string) error {
	conn := s.pool.Get()
	if conn == nil {
		return fmt.Errorf("failed to get connection from pool")
	}
	defer s.pool.Put(conn)

	// attach table on connection
	err := s.attachTableOnConnection(conn, tableAlias, table)
	if err != nil {
		return err
	}

	rows, err := conn.DBX.QueryxContext(ctx, fullQuery)
	if err != nil {
		if strings.Contains(err.Error(), "The provided token has expired.") ||
			strings.Contains(err.Error(), "AuthenticationFailed") ||
			strings.Contains(err.Error(), "Credential should be updated") {
			log.Warn("Credential error detected, attempting to refresh credential chain", zap.Error(err))
			err = s.refreshCredentialChainOnConnection(ctx, conn)
			if err != nil {
				return fmt.Errorf("failed to refresh credential chain after query failed: %w", err)
			}
			log.Info("Credential chain refreshed successfully, retrying query")
			rows, err = conn.DBX.QueryxContext(ctx, fullQuery)
			if err != nil {
				return fmt.Errorf("failed to query and scan with sqlx after credential refresh: %w", err)
			}
		} else if strings.Contains(err.Error(), "No files in log segment") {
			log.Warn("No files in log segment", zap.Error(err))
			return nil
		} else {
			return fmt.Errorf("failed to execute query: %w", err)
		}
	}
	defer rows.Close()

	destValue := reflect.ValueOf(dest)
	destSlice := destValue.Elem()

	// Get column names
	columns, err := rows.Columns()
	if err != nil {
		return fmt.Errorf("failed to get columns: %w", err)
	}

	log.Info("Processing map slice query", zap.Strings("columns", columns))

	// Process each row
	for rows.Next() {
		// Scan the row into interface{} slice
		values := make([]interface{}, len(columns))
		scanArgs := make([]interface{}, len(values))
		for i := range values {
			scanArgs[i] = &values[i]
		}

		err := rows.Scan(scanArgs...)
		if err != nil {
			log.Error("Failed to scan row", zap.Error(err))
			continue
		}

		// Create a new map for this row
		rowMap := make(map[string]interface{})
		for i, col := range columns {
			if i < len(values) {
				// Convert column name to lowercase to match expected field names
				colName := strings.ToLower(col)
				if colName == "binary_plan" {
					EnrichBinaryPlanFields(rowMap, binaryPlanRawString(values[i]))
				} else if colName == "plan" {
					if values[i] == nil {
						rowMap[colName] = nil
					} else {
						rowMap[colName] = DecodePlanField(binaryPlanRawString(values[i]))
					}
				} else if values[i] == nil {
					// Handle nil values - set to nil so JSON serializes as null
					rowMap[colName] = nil
				} else if strVal, ok := values[i].(string); ok && strVal == "null" {
					// Handle string "null" - convert to nil so JSON serializes as null instead of "null"
					rowMap[colName] = nil
				} else {
					rowMap[colName] = values[i]
				}
			}
		}

		// Append to the slice using reflection
		newSlice := reflect.Append(destSlice, reflect.ValueOf(rowMap))
		destSlice.Set(newSlice)
	}

	if err = rows.Err(); err != nil {
		return fmt.Errorf("error iterating rows: %w", err)
	}

	log.Info("Map slice query completed", zap.Int("rows", destSlice.Len()))
	return nil
}

func binaryPlanRawString(v interface{}) string {
	switch raw := v.(type) {
	case string:
		return raw
	case []byte:
		return string(raw)
	case nil:
		return ""
	default:
		return fmt.Sprint(raw)
	}
}

// selectIntoMapSlice handles scanning into []map[string]interface{}
// This function is kept for potential future use but currently unused
//
//nolint:unused
func (s *QueryDuckDBFetcher) selectIntoMapSlice(ctx context.Context, dest interface{}, fullQuery string) error {
	rows, err := s.dbx.QueryxContext(ctx, fullQuery)
	if err != nil {
		if strings.Contains(err.Error(), "The provided token has expired.") ||
			strings.Contains(err.Error(), "AuthenticationFailed") ||
			strings.Contains(err.Error(), "Credential should be updated") {
			log.Warn("Credential error detected, attempting to refresh credential chain", zap.Error(err))
			err = s.refreshCredentialChain(ctx)
			if err != nil {
				return fmt.Errorf("failed to refresh credential chain after query failed: %w", err)
			}
			log.Info("Credential chain refreshed successfully, retrying query")
			rows, err = s.dbx.QueryxContext(ctx, fullQuery)
			if err != nil {
				return fmt.Errorf("failed to query and scan with sqlx after credential refresh: %w", err)
			}
		} else if strings.Contains(err.Error(), "No files in log segment") {
			log.Warn("No files in log segment", zap.Error(err))
			return nil
		} else {
			return fmt.Errorf("failed to execute query: %w", err)
		}
	}
	defer rows.Close()

	destValue := reflect.ValueOf(dest)
	destSlice := destValue.Elem()

	// Get column names
	columns, err := rows.Columns()
	if err != nil {
		return fmt.Errorf("failed to get columns: %w", err)
	}

	log.Info("Processing map slice query", zap.Strings("columns", columns))

	// Process each row
	for rows.Next() {
		// Scan the row into interface{} slice
		values := make([]interface{}, len(columns))
		scanArgs := make([]interface{}, len(values))
		for i := range values {
			scanArgs[i] = &values[i]
		}

		err := rows.Scan(scanArgs...)
		if err != nil {
			log.Error("Failed to scan row", zap.Error(err))
			continue
		}

		// Create a new map for this row
		rowMap := make(map[string]interface{})
		for i, col := range columns {
			if i < len(values) {
				// Convert column name to lowercase to match expected field names
				colName := strings.ToLower(col)
				if colName == "binary_plan" {
					EnrichBinaryPlanFields(rowMap, binaryPlanRawString(values[i]))
				} else if colName == "plan" {
					if values[i] == nil {
						rowMap[colName] = nil
					} else {
						rowMap[colName] = DecodePlanField(binaryPlanRawString(values[i]))
					}
				} else if values[i] == nil {
					// Handle nil values - set to nil so JSON serializes as null
					rowMap[colName] = nil
				} else if strVal, ok := values[i].(string); ok && strVal == "null" {
					// Handle string "null" - convert to nil so JSON serializes as null instead of "null"
					rowMap[colName] = nil
				} else {
					rowMap[colName] = values[i]
				}
			}
		}

		// Append to the slice using reflection
		newSlice := reflect.Append(destSlice, reflect.ValueOf(rowMap))
		destSlice.Set(newSlice)
	}

	if err = rows.Err(); err != nil {
		return fmt.Errorf("error iterating rows: %w", err)
	}

	log.Info("Map slice query completed", zap.Int("rows", destSlice.Len()))
	return nil
}

func (s *QueryDuckDBFetcher) refreshCredentialChain(ctx context.Context) error {
	var refreshSQL string
	switch strings.ToLower(s.vendor) {
	case "aws":
		roleArn := os.Getenv("AWS_ROLE_ARN")
		credentials, err := getAWSAccessToken(s.region, roleArn)
		if err != nil {
			return fmt.Errorf("failed to refresh AWS access token: %w", err)
		}
		if credentials.SessionToken != "" {
			// With session token (for IAM roles, STS tokens, etc.)
			refreshSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					SESSION_TOKEN '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, credentials.SessionToken, s.region)
		} else {
			// Without session token (for IAM users)
			refreshSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, s.region)
		}
	case "alicloud":
		refreshSQL = fmt.Sprintf(aliReplaceSecret, s.region, s.region)
	case "gcp":
		refreshSQL = gcsReplaceSecret
	case "azure":
		// Azure blob storage uses access_code with Workload Identity token
		STORAGE_ACCOUNT := os.Getenv("AZURE_STORAGE_ACCOUNT")
		// Get new access token
		accessToken, err := getAzureAccessToken()
		if err != nil {
			return fmt.Errorf("failed to refresh Azure access token: %w", err)
		}
		refreshSQL = fmt.Sprintf(azureReplaceSecret, STORAGE_ACCOUNT, accessToken)
	case "local":
		return nil
	default:
		return fmt.Errorf("unsupported vendor for credential refresh: %s", s.vendor)
	}

	_, err := s.dbx.Exec(refreshSQL)
	if err != nil {
		log.Error("Failed to refresh credential chain secret", zap.Error(err), zap.String("vendor", s.vendor))
		return fmt.Errorf("failed to set secret: %w", err)
	}
	return nil
}

// attachTableOnConnection attaches a table on the specified connection (for connection pool)
func (s *QueryDuckDBFetcher) attachTableOnConnection(conn *DuckDBConnection, tableAlias, table string) error {
	if tableAlias == "" || table == "" {
		return nil
	}
	fullTablePath := table
	var disablePushdownFilters string
	// TODO: disable pushdown filters for slowlog table
	if false {
		disablePushdownFilters = ", PUSHDOWN_FILTERS 'none'"
	}
	_, err := conn.DBX.Exec(fmt.Sprintf("attach '%s' as %s(type delta%s);", fullTablePath, tableAlias, disablePushdownFilters))
	if err != nil {
		if strings.Contains(err.Error(), "Already attached") || strings.Contains(err.Error(), "already exists") {
			log.Debug("concurrent attach detected",
				zap.String("tableAlias", tableAlias),
				zap.String("table", table))
			return nil
		}
		return fmt.Errorf("failed to attach table: %w", err)
	}
	log.Info("table successfully attached on connection",
		zap.String("tableAlias", tableAlias),
		zap.String("attached", table))
	return nil
}

func (s *QueryDuckDBFetcher) attachTable(tableAlias, table string) error {
	if tableAlias == "" || table == "" {
		return nil
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if attachedTable, ok := s.attachedTables[tableAlias]; ok {
		if attachedTable != table {
			return fmt.Errorf("alias '%s' already attached to a different table '%s', cannot re-attach to '%s'", tableAlias, attachedTable, table)
		}
		log.Debug("table already attached, skipping re-attach.",
			zap.String("tableAlias", tableAlias),
			zap.String("attached", attachedTable))
		return nil
	} else {
		log.Info("attaching table.",
			zap.String("tableAlias", tableAlias),
			zap.String("table", table))

		fullTablePath := table
		var disablePushdownFilters string
		// TODO: disable pushdown filters for slowlog table
		if false {
			disablePushdownFilters = ", PUSHDOWN_FILTERS 'none'"
		}
		_, err := s.dbx.Exec(fmt.Sprintf("attach '%s' as %s(type delta%s);", fullTablePath, tableAlias, disablePushdownFilters))
		if err != nil {
			if strings.Contains(err.Error(), "Already attached") || strings.Contains(err.Error(), "already exists") {
				log.Info("concurrent attach detected",
					zap.String("tableAlias", tableAlias),
					zap.String("table", table))
				s.attachedTables[tableAlias] = table
				return nil
			}
			return fmt.Errorf("failed to attach table: %w", err)
		}
		s.attachedTables[tableAlias] = table
		log.Info("table successfully attached",
			zap.String("tableAlias", tableAlias),
			zap.String("attached", table))
	}
	return nil
}

// refreshCredentialChainOnConnection refreshes credentials on the specified connection (for connection pool)
func (s *QueryDuckDBFetcher) refreshCredentialChainOnConnection(ctx context.Context, conn *DuckDBConnection) error {
	var refreshSQL string
	switch strings.ToLower(s.vendor) {
	case "aws":
		roleArn := os.Getenv("AWS_ROLE_ARN")
		credentials, err := getAWSAccessToken(s.region, roleArn)
		if err != nil {
			return fmt.Errorf("failed to refresh AWS access token: %w", err)
		}
		if credentials.SessionToken != "" {
			// With session token (for IAM roles, STS tokens, etc.)
			refreshSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					SESSION_TOKEN '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, credentials.SessionToken, s.region)
		} else {
			// Without session token (for IAM users)
			refreshSQL = fmt.Sprintf(`
				CREATE OR REPLACE SECRET aws_secret (
					TYPE S3,
					KEY_ID '%s',
					SECRET '%s',
					REGION '%s'
				);
			`, credentials.AccessKeyID, credentials.SecretAccessKey, s.region)
		}
	case "alicloud":
		refreshSQL = fmt.Sprintf(aliReplaceSecret, s.region, s.region)
	case "gcp":
		refreshSQL = gcsReplaceSecret
	case "azure":
		// Azure blob storage uses access_code with Workload Identity token
		STORAGE_ACCOUNT := os.Getenv("AZURE_STORAGE_ACCOUNT")
		// Get new access token
		accessToken, err := getAzureAccessToken()
		if err != nil {
			return fmt.Errorf("failed to refresh Azure access token: %w", err)
		}
		refreshSQL = fmt.Sprintf(azureReplaceSecret, STORAGE_ACCOUNT, accessToken)
	case "local":
		return nil
	default:
		return fmt.Errorf("unsupported vendor for credential refresh: %s", s.vendor)
	}

	_, err := conn.DBX.Exec(refreshSQL)
	if err != nil {
		log.Error("Failed to refresh credential chain secret", zap.Error(err), zap.String("vendor", s.vendor))
		return fmt.Errorf("failed to set secret: %w", err)
	}
	return nil
}