package service

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/athena/types"
	"github.com/pingcap-inc/clinic/clinic_data_service/pkg/dto"
	"github.com/pingcap-inc/clinic/clinic_data_service/pkg/engine"
	"github.com/stretchr/testify/require"
)

type scriptedQueryEngine struct {
	errors   []error
	calls    int
	contexts []context.Context
}

func (e *scriptedQueryEngine) Execute(ctx context.Context, _ string, _ *dto.DataSourceInfo) (*dto.DataProxyQueryResult, error) {
	call := e.calls
	e.calls++
	e.contexts = append(e.contexts, ctx)
	if call < len(e.errors) && e.errors[call] != nil {
		return nil, e.errors[call]
	}
	return &dto.DataProxyQueryResult{}, nil
}

func (e *scriptedQueryEngine) Name() string                 { return "athena" }
func (e *scriptedQueryEngine) Health(context.Context) error { return nil }
func (e *scriptedQueryEngine) GetDatabase(string) string    { return "db" }
func (e *scriptedQueryEngine) GetTableSchema(context.Context, string, string, string) (*dto.SchemaResponse, error) {
	return nil, nil
}

type fakeAthenaRepairManager struct {
	repairCalls           int
	projectionRepairCalls int
	repairErr             error
	projectionRepairErr   error
	projectionRepairs     []engine.InstanceProjectionRepair
}

func (m *fakeAthenaRepairManager) EnsureTableExists(context.Context, string, string, string, string) error {
	return nil
}

func (m *fakeAthenaRepairManager) EnsureViewExists(context.Context, string, string, string, string, string, string) error {
	return nil
}

func (m *fakeAthenaRepairManager) IsViewTemplate(string) bool { return false }

func (m *fakeAthenaRepairManager) RepairDeltaSchemaMismatch(
	context.Context,
	[]engine.TableReference,
	string,
	string,
	string,
	time.Time,
) ([]string, error) {
	m.repairCalls++
	return []string{"db.slow_query_logs_42"}, m.repairErr
}

func (m *fakeAthenaRepairManager) RepairInstanceProjection(
	context.Context,
	[]engine.TableReference,
	string,
	string,
	string,
) ([]engine.InstanceProjectionRepair, error) {
	m.projectionRepairCalls++
	return m.projectionRepairs, m.projectionRepairErr
}

func invalidGlueSchemaError(queryID string) error {
	return &engine.AthenaQueryError{
		QueryID: queryID,
		State:   types.QueryExecutionStateFailed,
		Reason:  "INVALID_GLUE_SCHEMA: stale schema",
	}
}

func invalidInstanceProjectionError(queryID string) error {
	return &engine.AthenaQueryError{
		QueryID: queryID,
		State:   types.QueryExecutionStateFailed,
		Reason:  "INVALID_TABLE_PROPERTY: bad enum (property: projection.instance.values)",
	}
}

func TestExecuteQueryRepairsInstanceProjectionAndRetriesOnce(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{invalidInstanceProjectionError("first")}}
	manager := &fakeAthenaRepairManager{projectionRepairs: []engine.InstanceProjectionRepair{
		{Table: "topsql_tidb", InstanceCount: 2, Updated: true},
		{Table: "topsql_tikv", InstanceCount: 3, Updated: true},
	}}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{
		ClusterId: "42",
		Options: dto.DataProxyQueryOptions{
			DataSource: "s3://bucket/root",
		},
	}

	result, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT * FROM topsql_42",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		[]engine.TableReference{{Database: "db", Table: "topsql_42"}},
		req,
		time.Now(),
	)
	require.NoError(t, err)
	require.NotNil(t, result)
	require.Equal(t, 2, queryEngine.calls)
	require.Equal(t, 1, manager.projectionRepairCalls)
	require.Zero(t, manager.repairCalls)
}

func TestExecuteQueryDoesNotRetryInstanceProjectionTwice(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{
		invalidInstanceProjectionError("first"),
		invalidInstanceProjectionError("second"),
	}}
	manager := &fakeAthenaRepairManager{}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{Options: dto.DataProxyQueryOptions{DataSource: "s3://bucket/root"}}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		[]engine.TableReference{{Database: "db", Table: "topsql_tidb_42"}},
		req,
		time.Now(),
	)
	require.ErrorContains(t, err, "after Athena instance projection repair")
	require.Equal(t, 2, queryEngine.calls)
	require.Equal(t, 1, manager.projectionRepairCalls)
}

func TestExecuteQueryPreservesFirstFailureWhenProjectionRepairFails(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{invalidInstanceProjectionError("first")}}
	manager := &fakeAthenaRepairManager{projectionRepairErr: errors.New("S3 list failed")}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{Options: dto.DataProxyQueryOptions{DataSource: "s3://bucket/root"}}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		[]engine.TableReference{{Database: "db", Table: "topsql_tidb_42"}},
		req,
		time.Now(),
	)
	require.ErrorContains(t, err, "INVALID_TABLE_PROPERTY")
	require.ErrorContains(t, err, "Athena instance projection repair failed")
	require.ErrorContains(t, err, "S3 list failed")
	require.Equal(t, 1, queryEngine.calls)
}

func TestExecuteQueryRepairsAndRetriesOnce(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{invalidGlueSchemaError("first")}}
	manager := &fakeAthenaRepairManager{}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{
		ClusterId: "42",
		Options: dto.DataProxyQueryOptions{
			DataSource: "s3://bucket/root",
		},
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	result, err := svc.executeQueryWithAthenaSchemaRepair(
		ctx,
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		[]engine.TableReference{{Database: "db", Table: "slow_query_logs_42"}},
		req,
		time.Now(),
	)
	require.NoError(t, err)
	require.NotNil(t, result)
	require.Equal(t, 2, queryEngine.calls)
	require.Equal(t, 1, manager.repairCalls)
	require.Same(t, queryEngine.contexts[0], queryEngine.contexts[1])
}

func TestExecuteQueryDoesNotRepairOtherFailures(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{errors.New("syntax error")}}
	manager := &fakeAthenaRepairManager{}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"bad sql",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		nil,
		req,
		time.Now(),
	)
	require.Error(t, err)
	require.Equal(t, 1, queryEngine.calls)
	require.Zero(t, manager.repairCalls)
}

func TestExecuteQueryDoesNotRetryTwice(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{
		invalidGlueSchemaError("first"),
		invalidGlueSchemaError("second"),
	}}
	manager := &fakeAthenaRepairManager{}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{
		ClusterId: "42",
		Options: dto.DataProxyQueryOptions{
			DataSource: "s3://bucket/root",
		},
	}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		nil,
		req,
		time.Now(),
	)
	require.ErrorContains(t, err, "after Athena Delta schema repair")
	require.Equal(t, 2, queryEngine.calls)
	require.Equal(t, 1, manager.repairCalls)
}

func TestExecuteQueryDoesNotRetryWhenRepairFails(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{invalidGlueSchemaError("first")}}
	manager := &fakeAthenaRepairManager{repairErr: errors.New("Glue DeleteTable failed")}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{
		ClusterId: "42",
		Options: dto.DataProxyQueryOptions{
			DataSource: "s3://bucket/root",
		},
	}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		nil,
		req,
		time.Now(),
	)
	require.ErrorContains(t, err, "Athena Delta schema repair failed")
	require.Equal(t, 1, queryEngine.calls)
	require.Equal(t, 1, manager.repairCalls)
}

func TestExecuteQuerySkipsRepairWhenTableManagementDisabled(t *testing.T) {
	queryEngine := &scriptedQueryEngine{errors: []error{invalidGlueSchemaError("first")}}
	manager := &fakeAthenaRepairManager{}
	svc := &DataProxyService{tableManager: manager}
	req := &dto.DataProxyQueryRequest{
		Options: dto.DataProxyQueryOptions{SkipTableManagement: true},
	}

	_, err := svc.executeQueryWithAthenaSchemaRepair(
		context.Background(),
		queryEngine,
		"SELECT 1",
		&dto.DataSourceInfo{Vendor: "aws", Region: "us-west-2"},
		nil,
		req,
		time.Now(),
	)
	require.Error(t, err)
	require.Equal(t, 1, queryEngine.calls)
	require.Zero(t, manager.repairCalls)
}