package query

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"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/s3"
)

type athenaAPI interface {
	StartQueryExecution(ctx context.Context, params *athena.StartQueryExecutionInput, optFns ...func(*athena.Options)) (*athena.StartQueryExecutionOutput, error)
	GetQueryExecution(ctx context.Context, params *athena.GetQueryExecutionInput, optFns ...func(*athena.Options)) (*athena.GetQueryExecutionOutput, error)
	GetQueryResults(ctx context.Context, params *athena.GetQueryResultsInput, optFns ...func(*athena.Options)) (*athena.GetQueryResultsOutput, error)
	StopQueryExecution(ctx context.Context, params *athena.StopQueryExecutionInput, optFns ...func(*athena.Options)) (*athena.StopQueryExecutionOutput, error)
	GetWorkGroup(ctx context.Context, params *athena.GetWorkGroupInput, optFns ...func(*athena.Options)) (*athena.GetWorkGroupOutput, error)
	CreateWorkGroup(ctx context.Context, params *athena.CreateWorkGroupInput, optFns ...func(*athena.Options)) (*athena.CreateWorkGroupOutput, error)
}

type s3HeadAPI interface {
	HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
}

type athenaEngine struct {
	client   athenaAPI
	s3       s3HeadAPI
	presign  *s3.PresignClient
	cfg      Config
	catalog  string
	database string
}

func newAthenaEngine(cfg Config) (Engine, error) {
	awsCfg, err := config.LoadDefaultConfig(context.Background(), config.WithRegion(cfg.Region))
	if err != nil {
		return nil, fmt.Errorf("load aws config for athena: %w", err)
	}
	s3Client := s3.NewFromConfig(awsCfg)
	return &athenaEngine{
		client:   athena.NewFromConfig(awsCfg),
		s3:       s3Client,
		presign:  s3.NewPresignClient(s3Client),
		cfg:      cfg,
		catalog:  firstNonEmpty(cfg.AthenaCatalog, "AwsDataCatalog"),
		database: firstNonEmpty(cfg.AthenaDatabase, "o11y-logs"),
	}, nil
}

func (e *athenaEngine) Name() string { return "athena" }

func (e *athenaEngine) Submit(ctx context.Context, req *EngineSubmitRequest) (*Handle, error) {
	database := firstNonEmpty(req.Database, e.database)
	workgroup := firstNonEmpty(req.WorkGroup, e.cfg.ResolvedAthenaWorkGroup())
	output := normalizeS3URI(firstNonEmpty(req.OutputLocation, e.cfg.AthenaOutputLocation))

	in := &athena.StartQueryExecutionInput{
		QueryString: aws.String(req.SQL),
		QueryExecutionContext: &types.QueryExecutionContext{
			Catalog:  aws.String(e.catalog),
			Database: aws.String(database),
		},
	}
	if workgroup != "" {
		// Shared workgroup (o11y-olap) for billing; result S3 path is per-deployment.
		in.WorkGroup = aws.String(workgroup)
	}
	if output != "" {
		// Client-side ResultConfiguration so dedicated/serverless/ng can write to
		// their own bucket while sharing one workgroup. Requires the workgroup to
		// use classic S3 results with EnforceWorkGroupConfiguration=false.
		in.ResultConfiguration = &types.ResultConfiguration{
			OutputLocation: aws.String(output),
		}
	} else if workgroup == "" {
		return nil, fmt.Errorf("athena output location is required when workgroup is empty")
	}
	if req.ClientRequestToken != "" {
		in.ClientRequestToken = aws.String(req.ClientRequestToken)
	}

	out, err := e.client.StartQueryExecution(ctx, in)
	if err != nil {
		return nil, fmt.Errorf("athena StartQueryExecution: %w", err)
	}
	if out.QueryExecutionId == nil || *out.QueryExecutionId == "" {
		return nil, fmt.Errorf("athena returned empty QueryExecutionId")
	}
	return &Handle{EngineName: e.Name(), ExecutionID: *out.QueryExecutionId}, nil
}

func (e *athenaEngine) Status(ctx context.Context, handle *Handle) (*ExecutionStatus, error) {
	out, err := e.client.GetQueryExecution(ctx, &athena.GetQueryExecutionInput{
		QueryExecutionId: aws.String(handle.ExecutionID),
	})
	if err != nil {
		return nil, fmt.Errorf("athena GetQueryExecution: %w", err)
	}
	qe := out.QueryExecution
	if qe == nil || qe.Status == nil {
		return nil, fmt.Errorf("athena query execution missing status")
	}
	st := &ExecutionStatus{
		Status:            mapAthenaState(qe.Status.State),
		StateChangeReason: aws.ToString(qe.Status.StateChangeReason),
	}
	if qe.Status.SubmissionDateTime != nil {
		t := *qe.Status.SubmissionDateTime
		st.SubmittedAt = &t
	}
	if qe.Status.CompletionDateTime != nil {
		t := *qe.Status.CompletionDateTime
		st.CompletedAt = &t
	}
	if qe.Statistics != nil {
		st.ScannedBytes = aws.ToInt64(qe.Statistics.DataScannedInBytes)
		st.EngineExecutionTimeMillis = aws.ToInt64(qe.Statistics.EngineExecutionTimeInMillis)
		st.TotalExecutionTimeMillis = aws.ToInt64(qe.Statistics.TotalExecutionTimeInMillis)
		st.QueryQueueTimeMillis = aws.ToInt64(qe.Statistics.QueryQueueTimeInMillis)
		st.QueryPlanningTimeMillis = aws.ToInt64(qe.Statistics.QueryPlanningTimeInMillis)
	}
	return st, nil
}

func (e *athenaEngine) FetchPage(ctx context.Context, handle *Handle, nextToken string, maxRows int64) (*Page, error) {
	if maxRows <= 0 {
		maxRows = maxResultRows
	}
	if maxRows > maxResultRows {
		maxRows = maxResultRows
	}
	in := &athena.GetQueryResultsInput{
		QueryExecutionId: aws.String(handle.ExecutionID),
		MaxResults:       aws.Int32(int32(maxRows)),
	}
	if nextToken != "" {
		in.NextToken = aws.String(nextToken)
	}
	out, err := e.client.GetQueryResults(ctx, in)
	if err != nil {
		return nil, fmt.Errorf("athena GetQueryResults: %w", err)
	}
	page := &Page{}
	if out.NextToken != nil {
		page.NextToken = *out.NextToken
	}
	if out.ResultSet == nil {
		return page, nil
	}
	rs := out.ResultSet
	if rs.ResultSetMetadata != nil {
		for _, col := range rs.ResultSetMetadata.ColumnInfo {
			page.Columns = append(page.Columns, Column{
				Name: aws.ToString(col.Name),
				Type: aws.ToString(col.Type),
			})
		}
	}
	rows := rs.Rows
	// First page includes a header row matching column names.
	if nextToken == "" && len(rows) > 0 && isHeaderRow(rows[0], page.Columns) {
		rows = rows[1:]
	}
	for _, r := range rows {
		vals := make([]any, 0, len(r.Data))
		for _, d := range r.Data {
			if d.VarCharValue == nil {
				vals = append(vals, nil)
				continue
			}
			vals = append(vals, *d.VarCharValue)
		}
		page.Rows = append(page.Rows, vals)
	}
	return page, nil
}

func (e *athenaEngine) PresignResult(ctx context.Context, handle *Handle, expires time.Duration) (*PresignedResult, error) {
	if expires <= 0 {
		expires = defaultDownloadExpire
	}
	if expires > maxDownloadExpire {
		expires = maxDownloadExpire
	}
	out, err := e.client.GetQueryExecution(ctx, &athena.GetQueryExecutionInput{
		QueryExecutionId: aws.String(handle.ExecutionID),
	})
	if err != nil {
		return nil, fmt.Errorf("athena GetQueryExecution: %w", err)
	}
	qe := out.QueryExecution
	if qe == nil || qe.Status == nil {
		return nil, fmt.Errorf("athena query execution missing status")
	}
	if mapAthenaState(qe.Status.State) != StatusSucceeded {
		return nil, fmt.Errorf("query status is %s", mapAthenaState(qe.Status.State))
	}
	var outputLocation string
	if qe.ResultConfiguration != nil {
		outputLocation = aws.ToString(qe.ResultConfiguration.OutputLocation)
	}
	csvURI := resolveAthenaCSVLocation(outputLocation, handle.ExecutionID)
	if csvURI == "" {
		return nil, fmt.Errorf("athena result has no S3 output location (managed query results workgroups do not expose CSV in customer S3; use output_mode=inline or a workgroup with S3 ResultConfiguration)")
	}
	bucket, key, err := parseS3URI(csvURI)
	if err != nil {
		return nil, err
	}
	if e.presign == nil {
		return nil, fmt.Errorf("s3 presign client is not configured")
	}
	presigned, err := e.presign.PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	}, s3.WithPresignExpires(expires))
	if err != nil {
		return nil, fmt.Errorf("presign GetObject: %w", err)
	}
	result := &PresignedResult{
		Format:    "csv",
		S3URI:     csvURI,
		URL:       presigned.URL,
		ExpiresAt: time.Now().UTC().Add(expires),
	}
	if e.s3 != nil {
		if head, headErr := e.s3.HeadObject(ctx, &s3.HeadObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		}); headErr == nil && head.ContentLength != nil {
			n := *head.ContentLength
			result.Bytes = &n
		}
	}
	return result, nil
}

func (e *athenaEngine) Cancel(ctx context.Context, handle *Handle) error {
	_, err := e.client.StopQueryExecution(ctx, &athena.StopQueryExecutionInput{
		QueryExecutionId: aws.String(handle.ExecutionID),
	})
	if err != nil {
		return fmt.Errorf("athena StopQueryExecution: %w", err)
	}
	return nil
}

func mapAthenaState(state types.QueryExecutionState) string {
	switch state {
	case types.QueryExecutionStateQueued:
		return StatusQueued
	case types.QueryExecutionStateRunning:
		return StatusRunning
	case types.QueryExecutionStateSucceeded:
		return StatusSucceeded
	case types.QueryExecutionStateFailed:
		return StatusFailed
	case types.QueryExecutionStateCancelled:
		return StatusCancelled
	default:
		return string(state)
	}
}

func isHeaderRow(row types.Row, cols []Column) bool {
	if len(cols) == 0 || len(row.Data) != len(cols) {
		return false
	}
	for i, d := range row.Data {
		if d.VarCharValue == nil {
			return false
		}
		if !strings.EqualFold(*d.VarCharValue, cols[i].Name) {
			return false
		}
	}
	return true
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if strings.TrimSpace(v) != "" {
			return v
		}
	}
	return ""
}