package server

import (
	"context"
	"errors"
	"fmt"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	exportaction "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/export"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executorext"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

const tsoPhysicalShiftBits = 18

func (s *serverImpl) ListExports(ctx context.Context, clusterId string) ([]*entity.Export, error) {
	if clusterId == "" {
		tilog.Error(ctx, "cluster id is required when listing exports")
		return nil, status.Errorf(codes.InvalidArgument, "cluster id is required")
	}

	items, err := s.exportRepo.List(ctx, &repo.ExportListFilter{
		ClusterID: clusterId,
	})
	if err != nil {
		tilog.Error(ctx, "failed to list exports", zap.String("clusterId", clusterId), zap.Error(err))
		return nil, err
	}
	return items, nil
}

func (s *serverImpl) GetExport(ctx context.Context, clusterId, exportId string) (*entity.Export, error) {
	if exportId == "" {
		tilog.Error(ctx, "cluster id and export id are required when getting export")
		return nil, status.Errorf(codes.InvalidArgument, "export id is required")
	}
	exportItem, err := s.exportRepo.Load(ctx, exportId, true)
	if err != nil {
		tilog.Error(ctx, "failed to load export", zap.String("clusterId", clusterId), zap.String("exportId", exportId), zap.Error(err))
		return nil, err
	}
	return exportItem, nil
}

func (s *serverImpl) CreateExport(ctx context.Context, export *entity.Export) (*entity.Export, error) {
	if export == nil {
		tilog.Error(ctx, "export is required when creating export")
		return nil, status.Errorf(codes.InvalidArgument, "export is required")
	}
	if export.GetTidbGroupId() == "" {
		tilog.Error(ctx, "tidb group id is required when creating export")
		return nil, status.Errorf(codes.InvalidArgument, "tidb group id is required")
	}
	if snapshotTSO := export.GetSnapshotTso(); !isValidSnapshotTSO(snapshotTSO) {
		tilog.Error(ctx, "snapshot_tso has an invalid physical timestamp", zap.Uint64("snapshotTso", snapshotTSO))
		return nil, status.Errorf(codes.InvalidArgument, "snapshot_tso has an invalid physical timestamp")
	}
	if err := s.guardTiDBClusterMutationByID(ctx, export.GetClusterId(), sharedPoolMutationBlockedWhilePaused); err != nil {
		return nil, err
	}

	tilog.Info(ctx, "create export start")

	export.State = apiv1.Export_CREATING
	exportStr, err := protojson.Marshal(export)
	if err != nil {
		tilog.Error(ctx, "failed to marshal export when creating export",
			zap.String("clusterId", export.GetClusterId()),
			zap.String("exportId", export.GetExportId()),
			zap.Error(err))
		return nil, status.Errorf(codes.Internal, "failed to marshal export: %v", err)
	}

	var dagExecId string
	var dagSubmitted bool
	err = s.transactioner.Transaction(ctx, func(txCtx context.Context) error {
		if err := s.exportRepo.Create(txCtx, export); err != nil {
			if status.Code(err) == codes.AlreadyExists {
				tilog.Warn(ctx, "export already exists, skip create", zap.String("clusterId", export.GetClusterId()), zap.String("exportId", export.GetExportId()))
				return nil
			}
			tilog.Error(ctx, "failed to create export record",
				zap.String("clusterId", export.GetClusterId()),
				zap.String("exportId", export.GetExportId()),
				zap.Error(err))
			return err
		}

		if err := s.prepareTiDBGroupForExport(ctx, export); err != nil {
			tilog.Error(ctx, "failed to prepare tidb group for export",
				zap.String("clusterId", export.GetClusterId()),
				zap.String("exportId", export.GetExportId()),
				zap.String("tidbGroupId", export.GetTidbGroupId()),
				zap.Error(err))
			return err
		}

		dagArgs := actions.ActionArgs{
			workflow.ExportIdKey:      actions.ActionArgString(export.GetExportId()),
			workflow.ExportClusterKey: actions.ActionArgString(export.GetClusterId()),
			workflow.ExportKey:        actions.ActionArgString(string(exportStr)),
			utils.TraceIdKey:          utils.GetTraceId(ctx),
		}
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("create-export-%s", export.GetExportId()))
		if err := s.dagExecutor.Submit(txCtx, dag.PlanCreateExport, dagArgs, dagExecId); err != nil {
			tilog.Error(ctx, "failed to submit create export dag",
				zap.String("clusterId", export.GetClusterId()),
				zap.String("exportId", export.GetExportId()),
				zap.String("dagExecId", dagExecId),
				zap.Error(err))
			return status.Errorf(codes.Internal, "failed to submit create export dag: %v", err)
		}
		dagSubmitted = true
		return nil
	})
	if err != nil {
		if dagSubmitted {
			if err := s.dagExecutor.Abort(ctx, dagExecId); err != nil {
				tilog.Error(ctx, "failed to abort create export dag",
					zap.String("clusterId", export.GetClusterId()),
					zap.String("exportId", export.GetExportId()),
					zap.String("dagExecId", dagExecId),
					zap.Error(err))
			}
		}
		tilog.Error(ctx, "failed to create export",
			zap.String("clusterId", export.GetClusterId()),
			zap.String("exportId", export.GetExportId()),
			zap.Error(err))
		return nil, err
	}
	tilog.Info(ctx, "create export success", zap.String("exportId", export.GetExportId()), zap.String("dagExecId", dagExecId))
	return export, nil
}

func isValidSnapshotTSO(snapshotTSO uint64) bool {
	return snapshotTSO == 0 || snapshotTSO>>tsoPhysicalShiftBits != 0
}

func (s *serverImpl) prepareTiDBGroupForExport(ctx context.Context, export *entity.Export) error {
	cluster, err := s.tidbClusterRepo.Load(ctx, export.GetClusterId(), false)
	if err != nil {
		tilog.Error(ctx, "failed to load cluster when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.Error(err))
		return err
	}

	_, componentCfg, err := s.clusterConfigClient.GetExportComponentVersionAndConfig(ctx, cluster, apiv1.ComponentType_DUMPLING)
	if err != nil {
		tilog.Error(ctx, "failed to get export component config when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.Error(err))
		return err
	}
	if componentCfg == nil {
		return nil
	}

	dumplingCfg, err := exportaction.ParseExportDumplingConfig(componentCfg.GetConfigurationFileContent())
	if err != nil {
		tilog.Error(ctx, "failed to parse export dumpling config when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.Error(err))
		return err
	}

	tilog.Info(ctx, "got export dumpling config when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.Any("dumplingConfig", dumplingCfg))

	if !dumplingCfg.HasTiDBGroupScaleConfig() {
		tilog.Info(ctx, "skip scale because dumpling config for tidb group is not set", zap.String("exportId", export.GetExportId()))
		return nil
	}

	currentTiDBGroup, err := s.tidbGroupRepo.Load(ctx, export.GetTidbGroupId(), false)
	if err != nil {
		tilog.Error(ctx, "failed to load current tidb group when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.String("tidbGroupId", export.GetTidbGroupId()), zap.Error(err))
		return err
	}

	scaleTiDBGroup, shouldScale, err := buildExportTiDBGroupScaleTarget(ctx, export, currentTiDBGroup, dumplingCfg)
	if err != nil {
		tilog.Error(ctx, "failed to build export tidb group scale target when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.String("tidbGroupId", export.GetTidbGroupId()), zap.Error(err))
		return err
	}
	if !shouldScale {
		return nil
	}

	tilog.Info(ctx, "try to scale tidb group for export", zap.String("exportId", export.GetExportId()), zap.Any("nodeQuota", scaleTiDBGroup.NodeQuota))

	_, err = s.UpdateTiDBGroup(ctx, scaleTiDBGroup, true)
	if err != nil {
		tilog.Error(ctx, "failed to scale tidb group when preparing tidb group for export", zap.String("exportId", export.GetExportId()), zap.String("tidbGroupId", export.GetTidbGroupId()), zap.Error(err))
		return err
	}
	return nil
}

func buildExportTiDBGroupScaleTarget(
	ctx context.Context,
	export *entity.Export,
	currentTiDBGroup *entity.TiDBGroup,
	dumplingCfg *exportaction.ExportDumplingConfig,
) (*entity.TiDBGroup, bool, error) {
	shouldScale := false
	nodeQuota := currentTiDBGroup.GetNodeQuota()
	if nodeQuota == nil || nodeQuota.GetNodeSpec() == nil {
		tilog.Warn(ctx, "current tidb group has no node quota or node spec, cannot scale for export")
		return nil, false, nil
	}

	if dumplingCfg.TiDBGroupNodeCount > 0 && nodeQuota.GetNodeCount() != dumplingCfg.TiDBGroupNodeCount {
		nodeQuota.NodeCount = dumplingCfg.TiDBGroupNodeCount
		shouldScale = true
	}

	if dumplingCfg.TiDBGroupNodeSpec != "" {
		cpu, memoryGi, err := utils.ParseNodeSpec(dumplingCfg.TiDBGroupNodeSpec)
		if err != nil {
			return nil, false, fmt.Errorf("parse dumpling tidb group node spec: %w", err)
		}

		nodeSpec := nodeQuota.GetNodeSpec()
		if nodeSpec.Cpu != uint32(cpu) || nodeSpec.MemoryGi != uint32(memoryGi) {
			nodeQuota.NodeSpec = &entity.NodeSpec{
				Cpu:      uint32(cpu),
				MemoryGi: uint32(memoryGi),
			}
			shouldScale = true
		}
	}

	if !shouldScale {
		return nil, false, nil
	}

	return &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:        export.GetTidbGroupId(),
			ClusterId: export.GetClusterId(),
			NodeQuota: nodeQuota,
		},
	}, true, nil
}

func (s *serverImpl) CancelExport(ctx context.Context, clusterId, exportId string) (*entity.Export, error) {
	if exportId == "" {
		tilog.Error(ctx, "cluster id and export id are required when canceling export")
		return nil, status.Errorf(codes.InvalidArgument, "export id is required")
	}

	exportItem, err := s.exportRepo.Load(ctx, exportId, false)
	if err != nil {
		tilog.Error(ctx, "failed to load export when canceling",
			zap.String("clusterId", clusterId),
			zap.String("exportId", exportId),
			zap.Error(err))
		return nil, err
	}

	if exportItem.State == apiv1.Export_CANCELED {
		tilog.Warn(ctx, "export already canceled, skip cancel", zap.String("clusterId", clusterId), zap.String("exportId", exportId))
		return exportItem, nil
	}
	if err := s.guardTiDBClusterMutationByID(ctx, exportItem.GetClusterId(), sharedPoolMutationBlockedWhilePaused); err != nil {
		return nil, err
	}

	var dagExecId string
	var dagSubmitted bool
	err = s.transactioner.Transaction(ctx, func(txCtx context.Context) error {
		exportItem.State = apiv1.Export_CANCELED
		if err := s.exportRepo.Save(txCtx, exportItem); err != nil {
			tilog.Error(ctx, "failed to save export canceled state",
				zap.String("clusterId", clusterId),
				zap.String("exportId", exportId),
				zap.Error(err))
			return err
		}

		runningDag, err := s.dagInstanceFinder.GetRunningDagInstance(txCtx, dag.PlanCreateExport, executorext.DagInsFilter{
			ArgKey:    workflow.ExportIdKey,
			ArgValues: []string{exportId},
		})
		if err != nil && !errors.Is(err, executorext.DagInsNotFoundErr) {
			tilog.Error(ctx, "get running export failed", zap.String("exportId", exportId), zap.Error(err))
			return err
		}
		if runningDag != nil {
			tilog.Info(ctx, "abort running export dag", zap.String("exportId", exportId), zap.String("dagExecId", runningDag.ExecTask.ExecId))
			if err := s.dagExecutor.Abort(txCtx, runningDag.ExecTask.ExecId); err != nil {
				tilog.Error(ctx, "failed to abort create export dag", zap.String("exportId", exportId), zap.Error(err))
				return status.Errorf(codes.Internal, "failed to abort create export dag: %v", err)
			}
		}

		dagArgs := actions.ActionArgs{
			workflow.ExportIdKey:      actions.ActionArgString(exportItem.GetExportId()),
			workflow.ExportClusterKey: actions.ActionArgString(exportItem.GetClusterId()),
			utils.TraceIdKey:          utils.GetTraceId(ctx),
		}
		dagExecId = utils.GeneratedDagExecId(ctx, fmt.Sprintf("cancel-export-%s", exportItem.GetExportId()))
		if err := s.dagExecutor.Submit(txCtx, dag.PlanCancelExport, dagArgs, dagExecId); err != nil {
			tilog.Error(ctx, "failed to submit cancel export dag",
				zap.String("clusterId", clusterId),
				zap.String("exportId", exportId),
				zap.String("dagExecId", dagExecId),
				zap.Error(err))
			return status.Errorf(codes.Internal, "failed to submit cancel export dag: %v", err)
		}
		dagSubmitted = true
		return nil
	})
	if err != nil {
		if dagSubmitted {
			if err := s.dagExecutor.Abort(ctx, dagExecId); err != nil {
				tilog.Error(ctx, "failed to abort cancel export dag",
					zap.String("clusterId", clusterId),
					zap.String("exportId", exportId),
					zap.String("dagExecId", dagExecId),
					zap.Error(err))
			}
		}
		tilog.Error(ctx, "failed to cancel export", zap.String("clusterId", clusterId), zap.String("exportId", exportId), zap.Error(err))
		return nil, err
	}
	return exportItem, nil
}