package query

import (
	"context"
	"testing"

	"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/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type mockAthenaWorkGroupAPI struct {
	getOut    *athena.GetWorkGroupOutput
	getErr    error
	createIn  *athena.CreateWorkGroupInput
	createErr error
	createN   int
}

func (m *mockAthenaWorkGroupAPI) StartQueryExecution(context.Context, *athena.StartQueryExecutionInput, ...func(*athena.Options)) (*athena.StartQueryExecutionOutput, error) {
	return nil, nil
}
func (m *mockAthenaWorkGroupAPI) GetQueryExecution(context.Context, *athena.GetQueryExecutionInput, ...func(*athena.Options)) (*athena.GetQueryExecutionOutput, error) {
	return nil, nil
}
func (m *mockAthenaWorkGroupAPI) GetQueryResults(context.Context, *athena.GetQueryResultsInput, ...func(*athena.Options)) (*athena.GetQueryResultsOutput, error) {
	return nil, nil
}
func (m *mockAthenaWorkGroupAPI) StopQueryExecution(context.Context, *athena.StopQueryExecutionInput, ...func(*athena.Options)) (*athena.StopQueryExecutionOutput, error) {
	return nil, nil
}
func (m *mockAthenaWorkGroupAPI) GetWorkGroup(context.Context, *athena.GetWorkGroupInput, ...func(*athena.Options)) (*athena.GetWorkGroupOutput, error) {
	return m.getOut, m.getErr
}
func (m *mockAthenaWorkGroupAPI) CreateWorkGroup(_ context.Context, in *athena.CreateWorkGroupInput, _ ...func(*athena.Options)) (*athena.CreateWorkGroupOutput, error) {
	m.createN++
	m.createIn = in
	return &athena.CreateWorkGroupOutput{}, m.createErr
}

func TestDefaultConfigFromEnvWorkGroup(t *testing.T) {
	cfg := DefaultConfigFromEnv("aws", "us-east-1")
	assert.Empty(t, cfg.AthenaWorkGroup)
	assert.Empty(t, cfg.AthenaOutputLocation)

	cfg.ApplyLogsSettings("o11y-prod-shared-us-east-1-serverless", "")
	assert.Equal(t, "s3://o11y-prod-shared-us-east-1-serverless/athena-results/logs-query/", cfg.AthenaOutputLocation)
	assert.Equal(t, "o11y-olap-us-east-1-serverless", cfg.AthenaWorkGroup)
	assert.Equal(t, "us-east-1-serverless", cfg.LogsTableRegionKey())
}

func TestDefaultAthenaWorkGroupName(t *testing.T) {
	assert.Equal(t, "o11y-olap", defaultAthenaWorkGroupName(""))
	assert.Equal(t, "o11y-olap-us-west-2-staging", defaultAthenaWorkGroupName("us-west-2-staging"))
	assert.Equal(t, "o11y-olap-us-west-2", defaultAthenaWorkGroupName("us-west-2"))
}

func TestApplyLogsSettingsDoesNotOverrideExplicitOutput(t *testing.T) {
	cfg := DefaultConfigFromEnv("aws", "us-east-1")
	cfg.ApplyOverrides("", "", "", "s3://custom-bucket/results")
	cfg.ApplyLogsSettings("o11y-prod-shared-us-east-1-serverless", "")
	assert.Equal(t, "s3://custom-bucket/results/", cfg.AthenaOutputLocation)
	assert.Equal(t, "o11y-olap-us-east-1-serverless", cfg.AthenaWorkGroup)
}

func TestApplyLogsSettingsKeepsCustomWorkGroup(t *testing.T) {
	cfg := DefaultConfigFromEnv("aws", "us-east-1")
	cfg.ApplyOverrides("my-custom-wg", "", "", "")
	cfg.ApplyLogsSettings("o11y-prod-shared-us-east-1-serverless", "")
	assert.Equal(t, "my-custom-wg", cfg.AthenaWorkGroup)
}

func TestEnsureWorkGroupCreatesWhenMissing(t *testing.T) {
	mock := &mockAthenaWorkGroupAPI{
		getErr: &types.InvalidRequestException{Message: aws.String("WorkGroup o11y-olap-us-east-1-serverless is not found.")},
	}
	eng := &athenaEngine{
		client: mock,
		cfg: Config{
			LogsBucket:           "o11y-prod-shared-us-east-1-serverless",
			AthenaWorkGroup:      "o11y-olap-us-east-1-serverless",
			AthenaOutputLocation: "s3://o11y-prod-shared-us-east-1-serverless/athena-results/logs-query/",
		},
	}
	require.NoError(t, eng.EnsureWorkGroup(context.Background()))
	require.Equal(t, 1, mock.createN)
	require.NotNil(t, mock.createIn)
	assert.Equal(t, "o11y-olap-us-east-1-serverless", aws.ToString(mock.createIn.Name))
	require.NotNil(t, mock.createIn.Configuration)
	assert.False(t, aws.ToBool(mock.createIn.Configuration.EnforceWorkGroupConfiguration))
	require.NotNil(t, mock.createIn.Configuration.ResultConfiguration)
	assert.Equal(t, "s3://o11y-prod-shared-us-east-1-serverless/athena-results/logs-query/",
		aws.ToString(mock.createIn.Configuration.ResultConfiguration.OutputLocation))
	require.Len(t, mock.createIn.Tags, 1)
	assert.Equal(t, "project", aws.ToString(mock.createIn.Tags[0].Key))
	assert.Equal(t, "o11y-olap", aws.ToString(mock.createIn.Tags[0].Value))
}

func TestEnsureWorkGroupSkipsWhenExists(t *testing.T) {
	mock := &mockAthenaWorkGroupAPI{
		getOut: &athena.GetWorkGroupOutput{
			WorkGroup: &types.WorkGroup{
				Name: aws.String("o11y-olap-us-east-1-serverless"),
				Configuration: &types.WorkGroupConfiguration{
					ResultConfiguration: &types.ResultConfiguration{
						OutputLocation: aws.String("s3://o11y-prod-shared-us-east-1-serverless/athena-results/logs-query/"),
					},
				},
			},
		},
	}
	eng := &athenaEngine{
		client: mock,
		cfg: Config{
			AthenaWorkGroup:      "o11y-olap-us-east-1-serverless",
			AthenaOutputLocation: "s3://o11y-prod-shared-us-east-1-serverless/athena-results/logs-query/",
		},
	}
	require.NoError(t, eng.EnsureWorkGroup(context.Background()))
	assert.Equal(t, 0, mock.createN)
}

func TestEnsureWorkGroupRequiresOutput(t *testing.T) {
	eng := &athenaEngine{
		client: &mockAthenaWorkGroupAPI{},
		cfg:    Config{AthenaWorkGroup: "o11y-olap-us-east-1-serverless"},
	}
	err := eng.EnsureWorkGroup(context.Background())
	require.Error(t, err)
	assert.Contains(t, err.Error(), "output location is required")
}