package tidbtools

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"math"
	"strings"
	"testing"
	"time"

	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"github.com/tidbcloud/devops/pkg/infra/clients/coresvccli"
	"github.com/tidbcloud/devops/pkg/infra/clients/dataflowcli"
	nextgenglobal "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type premiumTestState struct {
	methods            []string
	backupSettingErr   error
	globalRestoreErr   error
	restoreProgressErr error
	restoreProgressMD  metadata.MD
	backupCatalogs     []*nextgenglobal.ListClusterBackupsResponse
	backupCatalogReq   []*nextgenglobal.ListClusterBackupsRequest
}

func newPremiumTestConn(t *testing.T, state *premiumTestState) *MockClientConnInterface {
	t.Helper()
	conn := NewMockClientConnInterface(gomock.NewController(t))
	conn.EXPECT().Invoke(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
		func(ctx context.Context, method string, request any, reply any, _ ...grpc.CallOption) error {
			return state.invoke(ctx, method, request, reply)
		},
	)
	return conn
}

func (c *premiumTestState) invoke(ctx context.Context, method string, request any, reply any) error {
	c.methods = append(c.methods, method)
	if c.backupSettingErr != nil && strings.HasSuffix(method, "/GetBackupSetting") {
		return c.backupSettingErr
	}
	if c.globalRestoreErr != nil && strings.Contains(method, "global.nextgen.v1.ClusterService/RestoreStatus") {
		return c.globalRestoreErr
	}
	if strings.HasSuffix(method, "/GetRestoreProgress") {
		c.restoreProgressMD, _ = metadata.FromOutgoingContext(ctx)
		if c.restoreProgressErr != nil {
			return c.restoreProgressErr
		}
	}
	if strings.HasSuffix(method, "/ListClusterBackups") {
		req, ok := request.(*nextgenglobal.ListClusterBackupsRequest)
		if !ok {
			return status.Error(codes.Internal, "unexpected backup catalog request")
		}
		c.backupCatalogReq = append(c.backupCatalogReq, req)
		out := reply.(*nextgenglobal.ListClusterBackupsResponse)
		if len(c.backupCatalogs) > 0 {
			*out = *c.backupCatalogs[0]
			c.backupCatalogs = c.backupCatalogs[1:]
		}
		return nil
	}
	switch out := reply.(type) {
	case *nextgenglobal.ListBackupsResponse:
		*out = nextgenglobal.ListBackupsResponse{}
	case *nextgenglobal.BackupSetting:
		*out = nextgenglobal.BackupSetting{}
	case *nextgenglobal.ListBackupPoliciesResponse:
		*out = nextgenglobal.ListBackupPoliciesResponse{}
	case *nextgenglobal.ListExportsResponse:
		*out = nextgenglobal.ListExportsResponse{}
	case *nextgenglobal.RestoreStatusResponse:
		*out = nextgenglobal.RestoreStatusResponse{
			State:    nextgenglobal.RestoreState_RESTORE_STATE_SUCCESS,
			Progress: 100,
			Details: &nextgenglobal.RestoreDetails{
				SourceClusterId:      41,
				RestoreMode:          nextgenglobal.RestoreMode_RESTORE_MODE_SNAPSHOT,
				BackupId:             72169,
				RestoreRequestedTime: timestamppb.New(time.Unix(100, 0)),
			},
		}
	default:
		return status.Errorf(codes.Unimplemented, "unexpected RPC %s", method)
	}
	return nil
}

func TestDedicatedRestoreProgressOnlyTreatsNotFoundAsEmpty(t *testing.T) {
	tests := []struct {
		name      string
		err       error
		wantEmpty bool
	}{
		{name: "not found", err: status.Error(codes.NotFound, "no active restore"), wantEmpty: true},
		{name: "internal", err: status.Error(codes.Internal, "dataflow failed")},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			state := &premiumTestState{restoreProgressErr: tt.err}
			conn := newPremiumTestConn(t, state)
			service := &svc{dataflow: &dataflowcli.DataflowCli{BackupSvc: dataflow.NewBackupServiceClient(conn)}}
			resp, err := service.GetRestoreProgress(context.Background(), 1, 2, 42)
			if tt.wantEmpty {
				if err != nil || resp != nil {
					t.Fatalf("expected empty restore progress, got resp=%#v err=%v", resp, err)
				}
				return
			}
			if status.Code(err) != codes.Internal || resp != nil {
				t.Fatalf("expected internal error, got resp=%#v err=%v", resp, err)
			}
		})
	}
}

func TestDedicatedRestoreProgressForwardsOpsUser(t *testing.T) {
	state := &premiumTestState{restoreProgressErr: status.Error(codes.NotFound, "no active restore")}
	conn := newPremiumTestConn(t, state)
	service := &svc{dataflow: &dataflowcli.DataflowCli{BackupSvc: dataflow.NewBackupServiceClient(conn)}}

	resp, err := service.GetRestoreProgress(context.Background(), 7, 8, 42)
	if err != nil || resp != nil {
		t.Fatalf("expected empty restore progress, got resp=%#v err=%v", resp, err)
	}

	values := state.restoreProgressMD.Get(dataflowOpsUserKey)
	if len(values) != 1 {
		t.Fatalf("expected one %q metadata value, got %#v", dataflowOpsUserKey, state.restoreProgressMD)
	}
	payload, err := base64.StdEncoding.DecodeString(values[0])
	if err != nil {
		t.Fatalf("decode operations-user metadata: %v", err)
	}
	var user dataflowOpsUser
	if err := json.Unmarshal(payload, &user); err != nil {
		t.Fatalf("unmarshal operations-user metadata: %v", err)
	}
	if user.TenantID != 7 || !user.IsImpersonated || len(user.Projects) != 1 || user.Projects[0].ID != 8 {
		t.Fatalf("unexpected operations user: %#v", user)
	}
}

func TestPremiumBackupSettingIsUnavailableDuringRestore(t *testing.T) {
	state := &premiumTestState{backupSettingErr: status.Error(codeImportClusterNotFound, "failed to get backup setting")}
	conn := newPremiumTestConn(t, state)
	service := &svc{
		coreSvc:       &coresvccli.CoreSvcCli{Mgmt: &coresvccli.MgmtService{}},
		nextGenGlobal: nextgenglobal.NewClusterServiceClient(conn),
	}

	setting, policies, err := service.GetPremiumBackupSetting(context.Background(), 1, 2, 42)
	if err != nil {
		t.Fatal(err)
	}
	if setting != nil || policies != nil {
		t.Fatalf("expected no setting or policies during restore, got setting=%#v policies=%#v", setting, policies)
	}
	if len(state.methods) != 1 || !strings.HasSuffix(state.methods[0], "/GetBackupSetting") {
		t.Fatalf("expected only GetBackupSetting to be called, got %v", state.methods)
	}
}

func TestPremiumBackupSettingPropagatesUnexpectedError(t *testing.T) {
	state := &premiumTestState{backupSettingErr: status.Error(codes.Internal, "management unavailable")}
	conn := newPremiumTestConn(t, state)
	service := &svc{
		coreSvc:       &coresvccli.CoreSvcCli{Mgmt: &coresvccli.MgmtService{}},
		nextGenGlobal: nextgenglobal.NewClusterServiceClient(conn),
	}

	_, _, err := service.GetPremiumBackupSetting(context.Background(), 1, 2, 42)
	if status.Code(err) != codes.Internal {
		t.Fatalf("expected internal error, got %v", err)
	}
}

func TestPremiumBackupRestoreUsesConsoleAndGlobalContracts(t *testing.T) {
	state := &premiumTestState{}
	conn := newPremiumTestConn(t, state)
	service := &svc{
		coreSvc: &coresvccli.CoreSvcCli{Mgmt: &coresvccli.MgmtService{
			NextGenExport: nextgenglobal.NewExportServiceClient(conn),
		}},
		nextGenGlobal: nextgenglobal.NewClusterServiceClient(conn),
		nextGenBackup: nextgenglobal.NewBackupServiceClient(conn),
	}

	ctx := context.Background()
	if _, err := service.ListPremiumBackups(ctx, 1, 2, 42, 1, 20); err != nil {
		t.Fatal(err)
	}
	if _, _, err := service.GetPremiumBackupSetting(ctx, 1, 2, 42); err != nil {
		t.Fatal(err)
	}
	if _, err := service.ListPremiumExports(ctx, 1, 2, 42, 1, 20); err != nil {
		t.Fatal(err)
	}
	if _, err := service.GetPremiumRestoreStatus(ctx, 1, 2, 42); err != nil {
		t.Fatal(err)
	}
	if len(state.methods) != 5 {
		t.Fatalf("expected 5 Premium API calls, got %d: %v", len(state.methods), state.methods)
	}
	if !strings.HasSuffix(state.methods[len(state.methods)-1], "/RestoreStatus") || !strings.Contains(state.methods[len(state.methods)-1], "global.nextgen.v1.ClusterService") {
		t.Fatalf("expected restore status to call ng-global, got %q", state.methods[len(state.methods)-1])
	}
}

func TestPremiumPageSkipRejectsOverflow(t *testing.T) {
	if _, err := premiumPageSkip(math.MaxInt32, 200); err == nil {
		t.Fatal("expected overflowing page offset to be rejected")
	}
	if skip, err := premiumPageSkip(2, 20); err != nil || skip != 20 {
		t.Fatalf("premiumPageSkip() = %d, %v", skip, err)
	}
}

func TestPremiumRestoreSourceFromGlobalDetails(t *testing.T) {
	details := &nextgenglobal.RestoreDetails{
		SourceClusterId:              41,
		RestoreMode:                  nextgenglobal.RestoreMode_RESTORE_MODE_PITR,
		BackupId:                     72169,
		PointInTime:                  timestamppb.New(time.Unix(90, 0)),
		RestoreRequestedTime:         timestamppb.New(time.Unix(100, 0)),
		RestoreUpdatedTime:           timestamppb.New(time.Unix(120, 0)),
		ExpectedRestoreDataSizeBytes: 1024,
		SourceRegionId:               "aws-us-west-2",
	}

	source := premiumRestoreSource(details)
	if source == nil || source.ClusterID != 41 || source.BackupID != "72169" || source.RestoreMode != "PITR" {
		t.Fatalf("unexpected restore source: %#v", source)
	}
	if source.PointInTime == nil || source.PointInTime.Unix() != 90 || source.EventAt == nil || source.EventAt.Unix() != 100 {
		t.Fatalf("unexpected restore timestamps: %#v", source)
	}
}

func TestPremiumRestoreStatusDoesNotFallBackWhileGlobalRecordIsMissing(t *testing.T) {
	state := &premiumTestState{globalRestoreErr: status.Error(codes.NotFound, "restore record not found")}
	conn := newPremiumTestConn(t, state)
	service := &svc{
		nextGenGlobal: nextgenglobal.NewClusterServiceClient(conn),
	}

	statusResp, err := service.GetPremiumRestoreStatus(context.Background(), 1, 2, 42)
	if err != nil {
		t.Fatal(err)
	}
	if statusResp != nil {
		t.Fatalf("expected no restore status, got %#v", statusResp)
	}
	if len(state.methods) != 1 || !strings.Contains(state.methods[0], "global.nextgen.v1.ClusterService") {
		t.Fatalf("expected only ng-global to be called, got %v", state.methods)
	}
}