package query

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/athena"
	"github.com/aws/aws-sdk-go-v2/service/athena/types"
	"github.com/aws/smithy-go"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

// EnsureWorkGroup checks that the shared Athena workgroup exists with classic
// S3 ResultConfiguration (not managed results). Creates it when missing.
// EnforceWorkGroupConfiguration is false so each deployment can pass its own
// output location while sharing one workgroup name for billing.
func (e *athenaEngine) EnsureWorkGroup(ctx context.Context) error {
	name := e.cfg.ResolvedAthenaWorkGroup()
	output := normalizeS3URI(e.cfg.AthenaOutputLocation)
	if output == "" {
		return fmt.Errorf("athena output location is required to ensure workgroup %q (set BUCKET_NAME or -athena-output-location)", name)
	}

	out, err := e.client.GetWorkGroup(ctx, &athena.GetWorkGroupInput{
		WorkGroup: aws.String(name),
	})
	if err == nil {
		return e.validateExistingWorkGroup(name, output, out.WorkGroup)
	}
	if !isAthenaWorkGroupNotFound(err) {
		return fmt.Errorf("get athena workgroup %q: %w", name, err)
	}

	log.Info("creating athena workgroup",
		zap.String("workgroup", name),
		zap.String("output_location", output))
	_, err = e.client.CreateWorkGroup(ctx, &athena.CreateWorkGroupInput{
		Name:        aws.String(name),
		Description: aws.String("o11y OLAP workgroup per logical region; per-query S3 output from client ResultConfiguration"),
		Tags: []types.Tag{
			{Key: aws.String(athenaWorkGroupProjectTagKey), Value: aws.String(athenaWorkGroupProjectTagValue)},
		},
		Configuration: &types.WorkGroupConfiguration{
			// Allow each logical-region deployment to set its own result bucket.
			EnforceWorkGroupConfiguration:   aws.Bool(false),
			PublishCloudWatchMetricsEnabled: aws.Bool(true),
			ResultConfiguration: &types.ResultConfiguration{
				OutputLocation: aws.String(output),
			},
		},
	})
	if err != nil {
		if isAthenaWorkGroupAlreadyExists(err) {
			log.Info("athena workgroup already exists (create race)", zap.String("workgroup", name))
			return nil
		}
		return fmt.Errorf("create athena workgroup %q: %w", name, err)
	}
	log.Info("athena workgroup created",
		zap.String("workgroup", name),
		zap.String("output_location", output))
	return nil
}

func (e *athenaEngine) validateExistingWorkGroup(name, wantOutput string, wg *types.WorkGroup) error {
	if wg == nil || wg.Configuration == nil || wg.Configuration.ResultConfiguration == nil {
		log.Warn("athena workgroup exists but has no S3 ResultConfiguration; export/presign may fail — recreate with customer S3 output or disable managed query results",
			zap.String("workgroup", name),
			zap.String("expected_output", wantOutput))
		return nil
	}
	got := normalizeS3URI(aws.ToString(wg.Configuration.ResultConfiguration.OutputLocation))
	if got == "" {
		log.Warn("athena workgroup ResultConfiguration.OutputLocation is empty",
			zap.String("workgroup", name),
			zap.String("expected_output", wantOutput))
		return nil
	}
	// Shared workgroup may keep a single default path; queries still pass
	// per-deployment ResultConfiguration when EnforceWorkGroupConfiguration=false.
	if aws.ToBool(wg.Configuration.EnforceWorkGroupConfiguration) {
		log.Warn("athena workgroup enforces its own settings; per-deployment result paths will be ignored — set EnforceWorkGroupConfiguration=false for shared o11y-olap billing workgroup",
			zap.String("workgroup", name),
			zap.String("workgroup_output", got))
	}
	log.Info("athena workgroup ready",
		zap.String("workgroup", name),
		zap.String("workgroup_output", got),
		zap.String("deployment_output", wantOutput))
	return nil
}

func isAthenaWorkGroupNotFound(err error) bool {
	var rnfe *types.ResourceNotFoundException
	if errors.As(err, &rnfe) {
		return true
	}
	var ire *types.InvalidRequestException
	if errors.As(err, &ire) {
		msg := strings.ToLower(ire.ErrorMessage())
		return strings.Contains(msg, "not found") || strings.Contains(msg, "does not exist")
	}
	var ae smithy.APIError
	if errors.As(err, &ae) {
		msg := strings.ToLower(ae.ErrorMessage())
		return ae.ErrorCode() == "InvalidRequestException" &&
			(strings.Contains(msg, "not found") || strings.Contains(msg, "does not exist"))
	}
	return false
}

func isAthenaWorkGroupAlreadyExists(err error) bool {
	var ire *types.InvalidRequestException
	if errors.As(err, &ire) {
		msg := strings.ToLower(ire.ErrorMessage())
		return strings.Contains(msg, "already exist")
	}
	var ae smithy.APIError
	if errors.As(err, &ae) {
		msg := strings.ToLower(ae.ErrorMessage())
		return strings.Contains(msg, "already exist")
	}
	return false
}