package engine

import (
	"context"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/pingcap-inc/clinic/clinic_data_service/pkg/dto"
	"github.com/pingcap-inc/clinic/clinic_data_service/pkg/log"
	"go.uber.org/zap"

	"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/athena"
	"github.com/aws/aws-sdk-go-v2/service/athena/types"
	"github.com/aws/aws-sdk-go-v2/service/glue"
	gluetypes "github.com/aws/aws-sdk-go-v2/service/glue/types"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/aws/smithy-go"
)

var athenaUnknownFieldPattern = regexp.MustCompile(`(?i)Column '([A-Za-z_][A-Za-z0-9_]*)' cannot be resolved`)

// AthenaEngine AWS Athena 查询引擎
type AthenaEngine struct {
	clients     sync.Map      // map[string]*athena.Client, key 为 region
	s3Clients   sync.Map      // map[string]*s3.Client, key 为 region
	glueClients sync.Map      // map[string]*glue.Client, key 为 region
	baseConfig  *AthenaConfig // 存储基础配置用于创建新客户端
	s3OutputLoc string        // 查询结果输出位置
	database    string        // 默认数据库
	workgroup   string        // 工作组
	mu          sync.RWMutex  // 保护客户端创建过程
}

// AthenaConfig Athena 引擎配置
type AthenaConfig struct {
	Region         string
	Database       string
	Workgroup      string
	OutputLocation string
	// 使用 AssumeRole 方式
	RoleArn string
	// 已废弃：使用 AK/SK 方式（保留用于向后兼容）
	AccessKey    string
	SecretKey    string
	SessionToken string
}

type s3ObjectLister interface {
	ListObjectsV2(context.Context, *s3.ListObjectsV2Input, ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
}

// AthenaQueryError preserves the structured failure information returned by
// GetQueryExecution so callers can make narrowly-scoped recovery decisions.
type AthenaQueryError struct {
	QueryID       string
	State         types.QueryExecutionState
	Reason        string
	ErrorCategory int32
	ErrorType     int32
	Retryable     bool
}

func (e *AthenaQueryError) Error() string {
	if e == nil {
		return ""
	}
	return fmt.Sprintf("query %s: %s", e.State, e.Reason)
}

// IsInvalidGlueSchema reports whether err is Athena's Delta Lake catalog
// mismatch. Athena exposes INVALID_GLUE_SCHEMA in StateChangeReason rather
// than as a dedicated SDK error code.
func IsInvalidGlueSchema(err error) bool {
	var queryErr *AthenaQueryError
	return errors.As(err, &queryErr) &&
		strings.Contains(queryErr.Reason, "INVALID_GLUE_SCHEMA:")
}

// IsInvalidInstanceProjection reports the narrowly scoped Athena table
// property failure produced by an empty or otherwise invalid instance enum.
func IsInvalidInstanceProjection(err error) bool {
	var queryErr *AthenaQueryError
	if !errors.As(err, &queryErr) {
		return false
	}
	reason := strings.ToLower(queryErr.Reason)
	return strings.Contains(reason, "invalid_table_property:") &&
		strings.Contains(reason, "property: projection.instance.values")
}

// AthenaQueryID returns the failed Athena query ID when available.
func AthenaQueryID(err error) string {
	var queryErr *AthenaQueryError
	if errors.As(err, &queryErr) {
		return queryErr.QueryID
	}
	return ""
}

func classifyAthenaUnknownField(err error) error {
	var queryErr *AthenaQueryError
	if !errors.As(err, &queryErr) || queryErr == nil {
		return nil
	}
	reason := queryErr.Reason
	if !strings.Contains(reason, "COLUMN_NOT_FOUND") &&
		!(strings.Contains(reason, "SYNTAX_ERROR") && athenaUnknownFieldPattern.MatchString(reason)) {
		return nil
	}
	matches := athenaUnknownFieldPattern.FindAllStringSubmatch(reason, -1)
	fields := make([]string, 0, len(matches))
	for _, match := range matches {
		if len(match) == 2 {
			fields = append(fields, match[1])
		}
	}
	return newUnknownFieldError("athena", queryErr.QueryID, fields)
}

// NewAthenaEngine 创建 Athena 引擎实例
func NewAthenaEngine(cfg *AthenaConfig) (*AthenaEngine, error) {
	engine := &AthenaEngine{
		baseConfig:  cfg,
		s3OutputLoc: cfg.OutputLocation,
		database:    cfg.Database,
		workgroup:   cfg.Workgroup,
	}

	// 预创建默认 region 的客户端
	defaultClient, err := engine.createClientForRegion(cfg.Region)
	if err != nil {
		return nil, fmt.Errorf("failed to create default client: %w", err)
	}
	engine.clients.Store(cfg.Region, defaultClient)

	return engine, nil
}

// getClientForRegion 获取或创建特定 region 的客户端
func (e *AthenaEngine) getClientForRegion(region string) (*athena.Client, error) {
	// 尝试从缓存获取
	if client, ok := e.clients.Load(region); ok {
		return client.(*athena.Client), nil
	}

	// 加锁创建新客户端
	e.mu.Lock()
	defer e.mu.Unlock()

	// 双重检查（避免重复创建）
	if client, ok := e.clients.Load(region); ok {
		return client.(*athena.Client), nil
	}

	// 创建并缓存新客户端
	client, err := e.createClientForRegion(region)
	if err != nil {
		return nil, err
	}

	e.clients.Store(region, client)
	return client, nil
}

// getS3ClientForRegion 获取或创建特定 region 的 S3 客户端
func (e *AthenaEngine) getS3ClientForRegion(region string) (*s3.Client, error) {
	// 尝试从缓存获取
	if client, ok := e.s3Clients.Load(region); ok {
		return client.(*s3.Client), nil
	}

	// 加锁创建新客户端
	e.mu.Lock()
	defer e.mu.Unlock()

	// 双重检查
	if client, ok := e.s3Clients.Load(region); ok {
		return client.(*s3.Client), nil
	}

	// 创建并缓存新客户端
	client, err := e.createS3ClientForRegion(region)
	if err != nil {
		return nil, err
	}

	e.s3Clients.Store(region, client)
	return client, nil
}

// HasDeltaLog reports whether dataSource contains a usable Delta transaction
// log. It uses the same region-aware, AssumeRole-backed S3 client as Athena.
// An entirely empty table root is considered unavailable rather than invalid.
func (e *AthenaEngine) HasDeltaLog(ctx context.Context, dataSource *dto.DataSourceInfo) (bool, error) {
	if dataSource == nil {
		return false, fmt.Errorf("data source is required")
	}

	region := dataSource.Region
	if region == "" && e.baseConfig != nil {
		region = e.baseConfig.Region
	}
	if region == "" {
		return false, fmt.Errorf("region is required to inspect Delta table availability")
	}

	client, err := e.getS3ClientForRegion(region)
	if err != nil {
		return false, err
	}

	return hasDeltaLog(ctx, client, dataSource)
}

func hasDeltaLog(ctx context.Context, client s3ObjectLister, dataSource *dto.DataSourceInfo) (bool, error) {
	if client == nil {
		return false, fmt.Errorf("S3 client is required")
	}
	if dataSource == nil {
		return false, fmt.Errorf("data source is required")
	}
	if dataSource.Bucket == "" {
		return false, fmt.Errorf("S3 bucket is required")
	}

	tableRoot := strings.Trim(dataSource.Path, "/")
	if tableRoot == "" {
		return false, fmt.Errorf("S3 table root is required")
	}
	tableRoot += "/"
	deltaLogPrefix := tableRoot + "_delta_log/"

	hasLogObject, err := hasNonDirectoryS3Object(ctx, client, dataSource.Bucket, deltaLogPrefix)
	if err != nil {
		return false, err
	}
	if hasLogObject {
		return true, nil
	}

	hasTableObject, err := hasNonDirectoryS3Object(ctx, client, dataSource.Bucket, tableRoot)
	if err != nil {
		return false, err
	}
	if !hasTableObject {
		return false, nil
	}

	return false, fmt.Errorf(
		"%w: s3://%s/%s contains objects but %s has no transaction-log objects",
		ErrInvalidDeltaSource,
		dataSource.Bucket,
		tableRoot,
		deltaLogPrefix,
	)
}

func hasNonDirectoryS3Object(ctx context.Context, client s3ObjectLister, bucket, prefix string) (bool, error) {
	var continuationToken *string
	for {
		output, err := client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
			Bucket:            aws.String(bucket),
			Prefix:            aws.String(prefix),
			ContinuationToken: continuationToken,
			MaxKeys:           aws.Int32(1000),
		})
		if err != nil {
			return false, err
		}

		for _, object := range output.Contents {
			key := aws.ToString(object.Key)
			if key != "" && !strings.HasSuffix(key, "/") {
				return true, nil
			}
		}

		if !aws.ToBool(output.IsTruncated) {
			return false, nil
		}
		if aws.ToString(output.NextContinuationToken) == "" {
			return false, fmt.Errorf("S3 ListObjectsV2 response for prefix %q is truncated without a continuation token", prefix)
		}
		continuationToken = output.NextContinuationToken
	}
}

// getGlueClientForRegion 获取或创建特定 region 的 Glue 客户端
func (e *AthenaEngine) getGlueClientForRegion(region string) (*glue.Client, error) {
	if client, ok := e.glueClients.Load(region); ok {
		return client.(*glue.Client), nil
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	if client, ok := e.glueClients.Load(region); ok {
		return client.(*glue.Client), nil
	}

	awsCfg, err := e.awsConfigForRegion(region)
	if err != nil {
		return nil, err
	}
	client := glue.NewFromConfig(awsCfg)
	e.glueClients.Store(region, client)
	return client, nil
}

// awsConfigForRegion 加载与 Athena/S3 一致的 AWS 配置（供 Glue 等复用）
func (e *AthenaEngine) awsConfigForRegion(region string) (aws.Config, error) {
	options := []func(*awsconfig.LoadOptions) error{
		awsconfig.WithRegion(region),
	}

	if e.baseConfig.RoleArn != "" {
		baseCfg, err := awsconfig.LoadDefaultConfig(context.Background(), awsconfig.WithRegion(region))
		if err != nil {
			return aws.Config{}, fmt.Errorf("failed to load base AWS config for region %s: %w", region, err)
		}
		stsClient := sts.NewFromConfig(baseCfg)
		credsProvider := stscreds.NewAssumeRoleProvider(stsClient, e.baseConfig.RoleArn)
		awsCfg, err := awsconfig.LoadDefaultConfig(context.Background(),
			awsconfig.WithRegion(region),
			awsconfig.WithCredentialsProvider(credsProvider),
		)
		if err != nil {
			return aws.Config{}, fmt.Errorf("failed to load AWS config with AssumeRole for region %s: %w", region, err)
		}
		return awsCfg, nil
	}

	if e.baseConfig.AccessKey != "" && e.baseConfig.SecretKey != "" {
		options = append(options, awsconfig.WithCredentialsProvider(
			aws.CredentialsProviderFunc(func(ctx context.Context) (aws.Credentials, error) {
				return aws.Credentials{
					AccessKeyID:     e.baseConfig.AccessKey,
					SecretAccessKey: e.baseConfig.SecretKey,
					SessionToken:    e.baseConfig.SessionToken,
				}, nil
			}),
		))
	}

	awsCfg, err := awsconfig.LoadDefaultConfig(context.Background(), options...)
	if err != nil {
		return aws.Config{}, fmt.Errorf("failed to load AWS config for region %s: %w", region, err)
	}
	return awsCfg, nil
}

// createS3ClientForRegion 创建特定 region 的 S3 客户端
func (e *AthenaEngine) createS3ClientForRegion(region string) (*s3.Client, error) {
	awsCfg, err := e.awsConfigForRegion(region)
	if err != nil {
		return nil, err
	}
	return s3.NewFromConfig(awsCfg), nil
}

// createClientForRegion 创建特定 region 的 Athena 客户端
func (e *AthenaEngine) createClientForRegion(region string) (*athena.Client, error) {
	awsCfg, err := e.awsConfigForRegion(region)
	if err != nil {
		return nil, err
	}
	return athena.NewFromConfig(awsCfg), nil
}

// GetGlueProjectionInstanceValues 从 Glue Data Catalog 读取表的 projection.instance.values（无则返回空字符串）
func (e *AthenaEngine) GetGlueProjectionInstanceValues(ctx context.Context, database, table, region string) (string, error) {
	params, err := e.GetGlueTableParameters(ctx, database, table, region)
	if err != nil {
		return "", err
	}
	return params["projection.instance.values"], nil
}

// GetGlueTableParameters reads Glue Data Catalog table parameters.
func (e *AthenaEngine) GetGlueTableParameters(ctx context.Context, database, table, region string) (map[string]string, error) {
	params, _, err := e.GetGlueTableInfo(ctx, database, table, region)
	return params, err
}

// GetGlueTableInfo reads the catalog properties needed to validate a managed
// Delta table before a destructive metadata rebuild.
func (e *AthenaEngine) GetGlueTableInfo(ctx context.Context, database, table, region string) (map[string]string, string, error) {
	gc, err := e.getGlueClientForRegion(region)
	if err != nil {
		return nil, "", err
	}
	out, err := gc.GetTable(ctx, &glue.GetTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(table),
	})
	if err != nil {
		return nil, "", err
	}
	if out.Table == nil {
		return nil, "", fmt.Errorf("glue GetTable returned empty table for %s.%s", database, table)
	}
	params := make(map[string]string, len(out.Table.Parameters))
	for k, v := range out.Table.Parameters {
		params[k] = v
	}
	location := ""
	if out.Table.StorageDescriptor != nil {
		location = aws.ToString(out.Table.StorageDescriptor.Location)
	}
	return params, location, nil
}

// DeleteGlueTable removes only the Glue Data Catalog table definition. Athena
// does not support DROP TABLE for Delta Lake tables, so schema recovery must
// delete the catalog entry through Glue before recreating it with Athena.
func (e *AthenaEngine) DeleteGlueTable(ctx context.Context, database, table, region string) error {
	gc, err := e.getGlueClientForRegion(region)
	if err != nil {
		return err
	}
	_, err = gc.DeleteTable(ctx, &glue.DeleteTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(table),
	})
	if err == nil {
		return nil
	}
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) && apiErr.ErrorCode() == "EntityNotFoundException" {
		return nil
	}
	return fmt.Errorf("glue DeleteTable failed: %w", err)
}

// UpdateGlueProjectionInstanceValues updates projection.instance.values directly in Glue Data Catalog.
func (e *AthenaEngine) UpdateGlueProjectionInstanceValues(ctx context.Context, database, table, region, values string) error {
	gc, err := e.getGlueClientForRegion(region)
	if err != nil {
		return err
	}

	out, err := gc.GetTable(ctx, &glue.GetTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(table),
	})
	if err != nil {
		return fmt.Errorf("glue GetTable failed: %w", err)
	}
	if out.Table == nil {
		return fmt.Errorf("glue GetTable returned empty table for %s.%s", database, table)
	}

	params := make(map[string]string, len(out.Table.Parameters)+1)
	for k, v := range out.Table.Parameters {
		params[k] = v
	}
	params["projection.instance.values"] = values

	skipArchive := true
	_, err = gc.UpdateTable(ctx, &glue.UpdateTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(table),
		SkipArchive:  aws.Bool(skipArchive),
		TableInput: &gluetypes.TableInput{
			Name:              out.Table.Name,
			Description:       out.Table.Description,
			LastAccessTime:    out.Table.LastAccessTime,
			LastAnalyzedTime:  out.Table.LastAnalyzedTime,
			Owner:             out.Table.Owner,
			Parameters:        params,
			PartitionKeys:     out.Table.PartitionKeys,
			Retention:         out.Table.Retention,
			StorageDescriptor: out.Table.StorageDescriptor,
			TableType:         out.Table.TableType,
			TargetTable:       out.Table.TargetTable,
			ViewExpandedText:  out.Table.ViewExpandedText,
			ViewOriginalText:  out.Table.ViewOriginalText,
		},
	})
	if err != nil {
		return fmt.Errorf("glue UpdateTable failed: %w", err)
	}
	return nil
}

// CheckDatabaseExists 通过 Glue Data Catalog 检查 database 是否存在
func (e *AthenaEngine) CheckDatabaseExists(ctx context.Context, database, region string) (bool, error) {
	gc, err := e.getGlueClientForRegion(region)
	if err != nil {
		return false, fmt.Errorf("failed to get Glue client for region %s: %w", region, err)
	}

	_, err = gc.GetDatabase(ctx, &glue.GetDatabaseInput{
		Name: aws.String(database),
	})
	if err != nil {
		var apiErr smithy.APIError
		if errors.As(err, &apiErr) && apiErr.ErrorCode() == "EntityNotFoundException" {
			return false, nil
		}
		if isCredentialExpiredError(err) {
			return false, fmt.Errorf("AWS credentials expired. Please refresh your access_key, secret_key, and session_token. Original error: %w", err)
		}
		return false, fmt.Errorf("glue GetDatabase failed: %w", err)
	}

	return true, nil
}

// Name 返回引擎名称
func (e *AthenaEngine) Name() string {
	return "athena"
}

// Execute 执行查询
func (e *AthenaEngine) Execute(ctx context.Context, sql string, dataSource *dto.DataSourceInfo) (*dto.DataProxyQueryResult, error) {
	return e.execute(ctx, sql, dataSource, true)
}

// ExecuteDDL executes Athena DDL and waits for completion without downloading
// its empty result file from S3.
func (e *AthenaEngine) ExecuteDDL(ctx context.Context, sql string, dataSource *dto.DataSourceInfo) error {
	_, err := e.execute(ctx, sql, dataSource, false)
	return err
}

func (e *AthenaEngine) execute(
	ctx context.Context,
	sql string,
	dataSource *dto.DataSourceInfo,
	fetchResult bool,
) (*dto.DataProxyQueryResult, error) {
	startTime := time.Now()

	// Region 必须由请求参数提供，避免跨 region 调用
	region := dataSource.Region
	if region == "" {
		// 如果配置文件有默认 region，使用它（向后兼容）
		if e.baseConfig.Region != "" {
			region = e.baseConfig.Region
			log.Warn(ctx, "using default region from config, please provide region in request",
				zap.String("region", region),
			)
		} else {
			return nil, fmt.Errorf("region is required in request parameters to avoid cross-region calls")
		}
	}

	// 获取或创建该 region 的客户端
	client, err := e.getClientForRegion(region)
	if err != nil {
		return nil, fmt.Errorf("failed to get client for region %s: %w", region, err)
	}

	// 1. 使用特定 region 的客户端启动查询
	startInput := &athena.StartQueryExecutionInput{
		QueryString: aws.String(sql),
		QueryExecutionContext: &types.QueryExecutionContext{
			Database: aws.String(e.database),
		},
	}

	// 动态生成输出位置，根据实际使用的 region
	// 只有在配置了输出位置时才设置 ResultConfiguration
	// 如果为空，则使用 Workgroup 的默认配置
	if e.s3OutputLoc != "" {
		// 将配置模板中的 {region} 占位符替换为实际请求的 region
		// 例如: s3://aws-athena-query-results-{region}-557537366020/ -> s3://aws-athena-query-results-ap-northeast-1-557537366020/
		outputLocation := strings.Replace(e.s3OutputLoc, "{region}", region, -1)
		startInput.ResultConfiguration = &types.ResultConfiguration{
			OutputLocation: aws.String(outputLocation),
		}
		log.Info(ctx, "Athena query output location configured",
			zap.String("region", region),
		)
	}

	// 如果指定了工作组，使用工作组
	if e.workgroup != "" {
		startInput.WorkGroup = aws.String(e.workgroup)
	}

	log.Info(ctx, "Starting Athena query execution",
		zap.String("database", e.database),
		zap.String("workgroup", e.workgroup),
		zap.String("region", region),
	)

	startOutput, err := client.StartQueryExecution(ctx, startInput)
	if err != nil {
		// 检测凭证过期错误
		if isCredentialExpiredError(err) {
			return nil, fmt.Errorf("AWS credentials expired. Please refresh your access_key, secret_key, and session_token. Original error: %w", err)
		}
		log.Error(ctx, "StartQueryExecution failed",
			zap.String("database", e.database),
			zap.String("workgroup", e.workgroup),
			zap.String("region", region),
		)
		return nil, fmt.Errorf("start query failed: %w", err)
	}

	queryID := *startOutput.QueryExecutionId

	// 2. 等待查询完成（传递 client）
	if err := e.waitForQuery(ctx, queryID, client); err != nil {
		if unknownErr := classifyAthenaUnknownField(err); unknownErr != nil {
			return nil, unknownErr
		}
		return nil, err
	}
	if !fetchResult {
		log.Info(ctx, "Athena DDL completed",
			zap.String("query_id", queryID),
			zap.String("region", region),
			zap.Duration("execution_time", time.Since(startTime)),
		)
		return nil, nil
	}

	// 3. 获取查询结果（从 S3 直接读取）
	result, err := e.fetchResults(ctx, queryID, client, region)
	if err != nil {
		return nil, err
	}

	// 4. 填充元数据（使用实际的 region）
	result.Metadata.ExecutionTime = dto.FormattedDuration(time.Since(startTime))
	result.Metadata.QueryID = queryID
	result.Metadata.Engine = e.Name()
	result.Metadata.Vendor = "aws"
	result.Metadata.Region = region

	return result, nil
}

// waitForQuery 等待查询完成
func (e *AthenaEngine) waitForQuery(ctx context.Context, queryID string, client *athena.Client) error {
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			getInput := &athena.GetQueryExecutionInput{
				QueryExecutionId: aws.String(queryID),
			}

			output, err := client.GetQueryExecution(ctx, getInput)
			if err != nil {
				return fmt.Errorf("get query execution failed: %w", err)
			}

			status := output.QueryExecution.Status.State
			switch status {
			case types.QueryExecutionStateSucceeded:
				return nil
			case types.QueryExecutionStateFailed, types.QueryExecutionStateCancelled:
				reason := ""
				if output.QueryExecution.Status.StateChangeReason != nil {
					reason = *output.QueryExecution.Status.StateChangeReason
				}
				queryErr := &AthenaQueryError{
					QueryID: queryID,
					State:   status,
					Reason:  reason,
				}
				if athenaErr := output.QueryExecution.Status.AthenaError; athenaErr != nil {
					queryErr.ErrorCategory = aws.ToInt32(athenaErr.ErrorCategory)
					queryErr.ErrorType = aws.ToInt32(athenaErr.ErrorType)
					queryErr.Retryable = athenaErr.Retryable
				}
				return queryErr
			}
			// 继续等待 QUEUED 或 RUNNING 状态
		}
	}
}

// fetchResults 获取查询结果（从 S3 直接读取）
func (e *AthenaEngine) fetchResults(ctx context.Context, queryID string, client *athena.Client, region string) (*dto.DataProxyQueryResult, error) {
	result := &dto.DataProxyQueryResult{
		Metadata: dto.DataProxyResultMetadata{
			QueryID: queryID,
		},
	}

	// 1. 获取查询执行信息，包含 S3 输出位置和引擎执行时间
	getExecInput := &athena.GetQueryExecutionInput{
		QueryExecutionId: aws.String(queryID),
	}
	execOutput, err := client.GetQueryExecution(ctx, getExecInput)
	if err != nil {
		return nil, fmt.Errorf("get query execution failed: %w", err)
	}

	// 填充引擎执行时间（Athena 后端统计）
	if execOutput.QueryExecution.Statistics != nil {
		if execOutput.QueryExecution.Statistics.EngineExecutionTimeInMillis != nil {
			result.Metadata.EngineExecutionTime = dto.FormattedDuration(time.Duration(*execOutput.QueryExecution.Statistics.EngineExecutionTimeInMillis) * time.Millisecond)
		}
		if execOutput.QueryExecution.Statistics.DataScannedInBytes != nil {
			result.Metadata.BytesScanned = *execOutput.QueryExecution.Statistics.DataScannedInBytes
		}
	}

	// 2. 提取 S3 结果文件路径
	s3OutputLocation := aws.ToString(execOutput.QueryExecution.ResultConfiguration.OutputLocation)
	if s3OutputLocation == "" {
		return nil, fmt.Errorf("no S3 output location found for query %s", queryID)
	}

	log.Info(ctx, "fetching results from S3",
		zap.String("query_id", queryID),
	)

	// 3. 解析 S3 URI (格式: s3://bucket/key)
	bucket, key, err := parseS3URI(s3OutputLocation)
	if err != nil {
		return nil, fmt.Errorf("parse S3 URI failed: %w", err)
	}
	log.Info(ctx, "S3 URI parsed",
		zap.String("query_id", queryID),
	)

	// 4. 获取 S3 客户端
	log.Info(ctx, "getting S3 client",
		zap.String("query_id", queryID),
		zap.String("region", region),
	)
	s3Client, err := e.getS3ClientForRegion(region)
	if err != nil {
		return nil, fmt.Errorf("failed to get S3 client: %w", err)
	}
	log.Info(ctx, "S3 client obtained",
		zap.String("query_id", queryID),
		zap.String("region", region),
	)

	// 5. 从 S3 下载结果文件
	log.Info(ctx, "downloading S3 object",
		zap.String("query_id", queryID),
	)
	getObjInput := &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	}
	getObjOutput, err := s3Client.GetObject(ctx, getObjInput)
	if err != nil {
		return nil, fmt.Errorf("get S3 object failed: %w", err)
	}
	defer getObjOutput.Body.Close()

	contentLength := int64(0)
	if getObjOutput.ContentLength != nil {
		contentLength = *getObjOutput.ContentLength
	}
	log.Info(ctx, "S3 object downloaded, starting to parse CSV",
		zap.String("query_id", queryID),
		zap.Int64("content_length", contentLength),
	)

	// 6. 解析 CSV 格式的结果
	csvReader := csv.NewReader(getObjOutput.Body)
	csvReader.LazyQuotes = true // 允许宽松的引号处理

	// 读取表头（第一行）
	header, err := csvReader.Read()
	if err != nil {
		if err == io.EOF {
			// 空结果文件（常见于 DDL 语句如 CREATE DATABASE, CREATE TABLE）
			log.Info(ctx, "empty result set from S3",
				zap.String("query_id", queryID),
			)
			result.Columns = []string{}
			result.Rows = [][]interface{}{}
			result.Metadata.RowCount = 0
			return result, nil
		}
		return nil, fmt.Errorf("read CSV header failed: %w", err)
	}
	result.Columns = header
	log.Info(ctx, "CSV header read",
		zap.String("query_id", queryID),
		zap.Int("column_count", len(header)),
	)

	// 读取数据行
	rowCount := int64(0)
	for {
		record, err := csvReader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("read CSV row failed: %w", err)
		}

		// 转换为 []interface{}
		var row []interface{}
		for _, field := range record {
			row = append(row, field)
		}
		result.Rows = append(result.Rows, row)
		rowCount++

		// 每 10000 行打印一次进度
		if rowCount%10000 == 0 {
			log.Info(ctx, "CSV parsing progress",
				zap.String("query_id", queryID),
				zap.Int64("rows_read", rowCount),
			)
		}
	}

	result.Metadata.RowCount = rowCount
	// S3 路径只在日志中记录，不再返回给客户端
	// result.Metadata.Result = s3OutputLocation (已删除)

	log.Info(ctx, "successfully fetched results from S3",
		zap.String("query_id", queryID),
		zap.Int64("row_count", rowCount),
	)

	return result, nil
}

// parseS3URI 解析 S3 URI，返回 bucket 和 key
func parseS3URI(s3URI string) (bucket, key string, err error) {
	// 格式: s3://bucket-name/path/to/file
	if !strings.HasPrefix(s3URI, "s3://") {
		return "", "", fmt.Errorf("invalid S3 URI format: %s", s3URI)
	}

	// 移除 s3:// 前缀
	path := strings.TrimPrefix(s3URI, "s3://")

	// 分割 bucket 和 key
	parts := strings.SplitN(path, "/", 2)
	if len(parts) != 2 {
		return "", "", fmt.Errorf("invalid S3 URI format: %s", s3URI)
	}

	return parts[0], parts[1], nil
}

// GetDatabase 根据 region 返回 Athena 默认数据库名（配置固定，region 可忽略）
func (e *AthenaEngine) GetDatabase(region string) string {
	return e.database
}

// Health 健康检查
func (e *AthenaEngine) Health(ctx context.Context) error {
	// 使用默认 region 的客户端
	client, err := e.getClientForRegion(e.baseConfig.Region)
	if err != nil {
		return fmt.Errorf("athena health check failed: %w", err)
	}

	_, err = client.ListWorkGroups(ctx, &athena.ListWorkGroupsInput{
		MaxResults: aws.Int32(1),
	})
	if err != nil {
		return fmt.Errorf("athena health check failed: %w", err)
	}
	return nil
}

// CheckTableExists 检查表是否存在
func (e *AthenaEngine) CheckTableExists(ctx context.Context, database, table, region string) (bool, error) {
	gc, err := e.getGlueClientForRegion(region)
	if err != nil {
		return false, fmt.Errorf("failed to get Glue client for region %s: %w", region, err)
	}

	_, err = gc.GetTable(ctx, &glue.GetTableInput{
		DatabaseName: aws.String(database),
		Name:         aws.String(table),
	})
	if err != nil {
		var apiErr smithy.APIError
		if errors.As(err, &apiErr) && apiErr.ErrorCode() == "EntityNotFoundException" {
			return false, nil
		}
		if isCredentialExpiredError(err) {
			return false, fmt.Errorf("AWS credentials expired. Please refresh your access_key, secret_key, and session_token. Original error: %w", err)
		}
		return false, fmt.Errorf("glue GetTable failed: %w", err)
	}

	return true, nil
}

// GetTableSchema 获取表结构
func (e *AthenaEngine) GetTableSchema(ctx context.Context, database, table, region string) (*dto.SchemaResponse, error) {
	// 执行 DESCRIBE TABLE 查询
	// 使用反引号以支持包含特殊字符（如连字符）的数据库名
	sql := fmt.Sprintf("DESCRIBE `%s`.`%s`", database, table)

	client, err := e.getClientForRegion(region)
	if err != nil {
		return nil, fmt.Errorf("failed to get client: %w", err)
	}

	// 启动查询
	startInput := &athena.StartQueryExecutionInput{
		QueryString: aws.String(sql),
		QueryExecutionContext: &types.QueryExecutionContext{
			Database: aws.String(e.database),
		},
	}

	if e.s3OutputLoc != "" {
		// 替换 {region} 占位符，与主查询保持一致
		outputLocation := strings.Replace(e.s3OutputLoc, "{region}", region, -1)
		startInput.ResultConfiguration = &types.ResultConfiguration{
			OutputLocation: aws.String(outputLocation),
		}
	}

	if e.workgroup != "" {
		startInput.WorkGroup = aws.String(e.workgroup)
	}

	startOutput, err := client.StartQueryExecution(ctx, startInput)
	if err != nil {
		// 检测凭证过期错误
		if isCredentialExpiredError(err) {
			return nil, fmt.Errorf("AWS credentials expired. Please refresh your access_key, secret_key, and session_token. Original error: %w", err)
		}
		return nil, fmt.Errorf("start query failed: %w", err)
	}

	queryID := *startOutput.QueryExecutionId

	// 等待查询完成
	if err := e.waitForQuery(ctx, queryID, client); err != nil {
		return nil, fmt.Errorf("wait for query failed: %w", err)
	}

	// 获取结果
	result, err := e.fetchResults(ctx, queryID, client, region)
	if err != nil {
		return nil, fmt.Errorf("fetch results failed: %w", err)
	}

	// 解析结果为 ColumnInfo
	// 注意：Athena 的 DESCRIBE 返回的是 tab 分隔的单个字符串
	// 格式：普通列列表 -> "# Partition Information" -> 分区列列表
	var columns []dto.ColumnInfo
	var partitions []dto.PartitionDef
	inPartitionSection := false

	for _, row := range result.Rows {
		if len(row) == 0 {
			continue
		}

		// Athena DESCRIBE 返回格式：每行是一个 tab 分隔的字符串
		// 格式："col_name\tcol_type\tcomment"
		line, ok := row[0].(string)
		if !ok {
			continue
		}

		// 检查是否到达分区信息段
		if strings.HasPrefix(strings.TrimSpace(line), "#") {
			if strings.Contains(line, "Partition") {
				inPartitionSection = true
			}
			continue
		}

		// 按 tab 分割
		parts := strings.Split(line, "\t")
		if len(parts) < 2 {
			continue
		}

		// 清理空格
		colName := strings.TrimSpace(parts[0])
		colType := strings.TrimSpace(parts[1])

		if colName == "" || colName == "col_name" {
			continue
		}

		// 根据所在段决定是普通列还是分区列
		if inPartitionSection {
			partitions = append(partitions, dto.PartitionDef{
				Name: colName,
				Type: colType,
			})
		} else {
			column := dto.ColumnInfo{
				Name: colName,
				Type: colType,
			}

			// 如果有第3列，视为注释
			if len(parts) >= 3 {
				comment := strings.TrimSpace(parts[2])
				if comment != "" {
					column.Comment = comment
				}
			}

			columns = append(columns, column)
		}
	}

	// 从 columns 中移除所有分区列
	partitionNames := make(map[string]bool)
	for _, p := range partitions {
		partitionNames[p.Name] = true
	}

	var filteredColumns []dto.ColumnInfo
	for _, col := range columns {
		if !partitionNames[col.Name] {
			filteredColumns = append(filteredColumns, col)
		}
	}

	// 获取表的 LOCATION 和类型（通过 SHOW CREATE TABLE）
	location, isView, err := e.getTableLocationAndType(ctx, database, table, region)
	if err != nil {
		// 获取 location 失败不影响返回 schema
		location = ""
	}

	return &dto.SchemaResponse{
		Database:   database,
		Table:      table,
		Columns:    filteredColumns,
		Partitions: partitions,
		Location:   location,
		IsView:     isView,
	}, nil
}

// getTableLocationAndType 获取表的 S3 存储位置和类型（表/VIEW）
func (e *AthenaEngine) getTableLocationAndType(ctx context.Context, database, table, region string) (string, bool, error) {
	// 使用反引号支持包含特殊字符的数据库名
	sql := fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`", database, table)

	client, err := e.getClientForRegion(region)
	if err != nil {
		return "", false, fmt.Errorf("failed to get client: %w", err)
	}

	startInput := &athena.StartQueryExecutionInput{
		QueryString: aws.String(sql),
		QueryExecutionContext: &types.QueryExecutionContext{
			Database: aws.String(e.database),
		},
	}

	if e.s3OutputLoc != "" {
		// 替换 {region} 占位符，与主查询保持一致
		outputLocation := strings.Replace(e.s3OutputLoc, "{region}", region, -1)
		startInput.ResultConfiguration = &types.ResultConfiguration{
			OutputLocation: aws.String(outputLocation),
		}
	}

	if e.workgroup != "" {
		startInput.WorkGroup = aws.String(e.workgroup)
	}

	startOutput, err := client.StartQueryExecution(ctx, startInput)
	if err != nil {
		// 检测凭证过期错误
		if isCredentialExpiredError(err) {
			return "", false, fmt.Errorf("AWS credentials expired. Please refresh your access_key, secret_key, and session_token. Original error: %w", err)
		}
		return "", false, fmt.Errorf("start query failed: %w", err)
	}

	queryID := *startOutput.QueryExecutionId
	if err := e.waitForQuery(ctx, queryID, client); err != nil {
		return "", false, fmt.Errorf("wait for query failed: %w", err)
	}

	result, err := e.fetchResults(ctx, queryID, client, region)
	if err != nil {
		return "", false, fmt.Errorf("fetch results failed: %w", err)
	}

	// 解析 CREATE TABLE/VIEW 语句，提取 LOCATION 和类型
	// 注意：SHOW CREATE TABLE 返回的每行是语句的一部分，需要拼接
	var fullStatement strings.Builder
	for _, row := range result.Rows {
		if len(row) > 0 {
			if line, ok := row[0].(string); ok {
				fullStatement.WriteString(line)
				fullStatement.WriteString(" ")
			}
		}
	}

	createStmt := fullStatement.String()
	isView := strings.Contains(createStmt, "CREATE VIEW") || strings.Contains(createStmt, "CREATE OR REPLACE VIEW")

	// VIEW 没有 LOCATION
	if isView {
		return "", true, nil
	}

	// 简单正则匹配 LOCATION 's3://...'
	if strings.Contains(createStmt, "LOCATION") {
		start := strings.Index(createStmt, "LOCATION")
		if start != -1 {
			locationPart := createStmt[start:]
			// 提取引号内的内容
			if idx1 := strings.Index(locationPart, "'"); idx1 != -1 {
				if idx2 := strings.Index(locationPart[idx1+1:], "'"); idx2 != -1 {
					location := locationPart[idx1+1 : idx1+1+idx2]
					return location, false, nil
				}
			}
		}
	}

	return "", false, fmt.Errorf("location not found in CREATE TABLE statement")
}

// ListS3Prefixes 列出 S3 路径下的所有子前缀（使用 Delimiter="/" 模拟 ls 目录）
// 返回 CommonPrefixes 中的前缀列表（不含 inputPrefix 本身）
func (e *AthenaEngine) ListS3Prefixes(ctx context.Context, bucket, prefix, region string) ([]string, error) {
	s3Client, err := e.getS3ClientForRegion(region)
	if err != nil {
		return nil, fmt.Errorf("failed to get S3 client for region %s: %w", region, err)
	}

	delimiter := "/"
	input := &s3.ListObjectsV2Input{
		Bucket:    aws.String(bucket),
		Prefix:    aws.String(prefix),
		Delimiter: aws.String(delimiter),
	}

	var prefixes []string
	paginator := s3.NewListObjectsV2Paginator(s3Client, input)
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, fmt.Errorf("failed to list S3 prefixes: %w", err)
		}
		for _, cp := range page.CommonPrefixes {
			if cp.Prefix != nil {
				prefixes = append(prefixes, *cp.Prefix)
			}
		}
	}

	return prefixes, nil
}

// isCredentialExpiredError 检查是否是凭证过期或无效错误
func isCredentialExpiredError(err error) bool {
	if err == nil {
		return false
	}
	errMsg := strings.ToLower(err.Error())
	return strings.Contains(errMsg, "expiredtoken") ||
		strings.Contains(errMsg, "invalidclienttokenid") ||
		strings.Contains(errMsg, "securitytokenexpired") ||
		strings.Contains(errMsg, "token has expired") ||
		strings.Contains(errMsg, "credentials have expired") ||
		strings.Contains(errMsg, "unrecognizedclientexception") ||
		strings.Contains(errMsg, "security token included in the request is invalid")
}

// fetchResults 获取查询结果
func (e *AthenaEngine) fetchResults_bak(ctx context.Context, queryID string, client *athena.Client) (*dto.DataProxyQueryResult, error) {
	result := &dto.DataProxyQueryResult{
		Metadata: dto.DataProxyResultMetadata{
			QueryID: queryID,
		},
	}

	var nextToken *string
	isFirstPage := true

	// 循环获取所有分页结果
	for {
		input := &athena.GetQueryResultsInput{
			QueryExecutionId: aws.String(queryID),
			MaxResults:       aws.Int32(1000), // 每次最多获取 1000 行
			NextToken:        nextToken,
		}

		output, err := client.GetQueryResults(ctx, input)
		log.Info(ctx, "get query results",
			zap.String("query_id", queryID),
			zap.String("next_token", aws.ToString(nextToken)),
			zap.Int32("max_results", *input.MaxResults),
		)
		if err != nil {
			return nil, fmt.Errorf("get query results failed: %w", err)
		}

		// 解析结果
		if len(output.ResultSet.Rows) == 0 {
			break
		}

		// 第一页：提取列名（第一行）
		startIndex := 0
		if isFirstPage {
			headerRow := output.ResultSet.Rows[0]
			for _, col := range headerRow.Data {
				if col.VarCharValue != nil {
					result.Columns = append(result.Columns, *col.VarCharValue)
				}
			}
			startIndex = 1 // 跳过表头
			isFirstPage = false
		}

		// 提取数据行
		for _, row := range output.ResultSet.Rows[startIndex:] {
			var rowData []interface{}
			for _, col := range row.Data {
				if col.VarCharValue != nil {
					rowData = append(rowData, *col.VarCharValue)
				} else {
					rowData = append(rowData, nil)
				}
			}
			result.Rows = append(result.Rows, rowData)
		}

		// 检查是否有下一页
		if output.NextToken == nil {
			break
		}
		nextToken = output.NextToken
	}

	result.Metadata.RowCount = int64(len(result.Rows))

	return result, nil
}