package nodegroup

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

	"github.com/google/uuid"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	ngspec "github.com/tidbcloud/cluster-service-ng/pkg/nodegroup"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/pkgv2/gormutil"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gorm.io/gorm"
)

type stateRepoImpl struct {
	provider db.Provider
}

func normalizeNodeSpecForOutput(spec string) string {
	normalized := normalizeNodeSpec(spec)
	cpu, mem, err := ngspec.ParseSpec(normalized)
	if err != nil {
		return strings.ToUpper(strings.TrimSpace(spec))
	}
	return fmt.Sprintf("%dC%dG", cpu, mem)
}

// NewDataplaneNodeInstanceGroupStateRepo creates a new DataplaneNodeInstanceGroupStateRepo.
func NewDataplaneNodeInstanceGroupStateRepo(provider db.Provider) repo.DataplaneNodeInstanceGroupStateRepo {
	return &stateRepoImpl{provider: provider}
}

func (r *stateRepoImpl) Save(ctx context.Context, state *entity.DataplaneNodeInstanceGroupState) error {
	if state == nil {
		return status.Errorf(codes.InvalidArgument, "state is nil")
	}

	wdb := r.provider.WriteDB(ctx)

	// Generate ID if not set
	if state.ID == "" {
		state.ID = uuid.New().String()
	}

	model := &schema.DataplaneNodeInstanceGroupState{
		ID:                 state.ID,
		DataplaneName:      state.DataplaneName,
		NodeSpec:           normalizeNodeSpec(state.NodeSpec),
		PrimaryInstance:    state.PrimaryInstance,
		State:              state.State,
		StateEnteredAt:     state.StateEnteredAt,
		FencingToken:       state.FencingToken,
		LastShortageReason: state.LastShortageReason,
		CooldownUntil:      state.CooldownUntil,
	}

	if err := wdb.Save(model).Error; err != nil {
		return status.Errorf(codes.Internal, "failed to save DataplaneNodeInstanceGroupState: %v", err)
	}

	return nil
}

func (r *stateRepoImpl) Load(ctx context.Context, id string) (*entity.DataplaneNodeInstanceGroupState, error) {
	if id == "" {
		return nil, status.Errorf(codes.InvalidArgument, "id is empty")
	}

	rdb := r.provider.ReadDB(ctx)

	var model schema.DataplaneNodeInstanceGroupState
	if err := rdb.Where("id = ?", id).First(&model).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, status.Errorf(codes.NotFound, "NodeInstanceGroupState not found: %s", id)
		}
		return nil, status.Errorf(codes.Internal, "failed to load DataplaneNodeInstanceGroupState: %v", err)
	}
	model.NodeSpec = normalizeNodeSpecForOutput(model.NodeSpec)

	return &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: model}, nil
}

func (r *stateRepoImpl) LoadByDataplaneAndSpec(ctx context.Context, dataplaneName, nodeSpec, primaryInstance string) (*entity.DataplaneNodeInstanceGroupState, error) {
	if dataplaneName == "" || nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dataplaneName, nodeSpec, and primaryInstance are required")
	}
	nodeSpec = normalizeNodeSpec(nodeSpec)

	rdb := r.provider.ReadDB(ctx)

	var model schema.DataplaneNodeInstanceGroupState
	if err := rdb.Where("dataplane_name = ? AND node_spec = ? AND primary_instance = ?", dataplaneName, nodeSpec, primaryInstance).First(&model).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, nil //nolint:nilnil // Interface contract: returns nil without error when not found
		}
		return nil, status.Errorf(codes.Internal, "failed to load DataplaneNodeInstanceGroupState: %v", err)
	}
	model.NodeSpec = normalizeNodeSpecForOutput(model.NodeSpec)

	return &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: model}, nil
}

// loadByDataplaneAndSpecUnscoped is like LoadByDataplaneAndSpec but includes soft-deleted records.
func (r *stateRepoImpl) loadByDataplaneAndSpecUnscoped(ctx context.Context, dataplaneName, nodeSpec, primaryInstance string) (*schema.DataplaneNodeInstanceGroupState, error) {
	nodeSpec = normalizeNodeSpec(nodeSpec)
	rdb := r.provider.ReadDB(ctx)

	var model schema.DataplaneNodeInstanceGroupState
	if err := rdb.Unscoped().Where("dataplane_name = ? AND node_spec = ? AND primary_instance = ?", dataplaneName, nodeSpec, primaryInstance).First(&model).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, nil //nolint:nilnil // returns nil when not found
		}
		return nil, status.Errorf(codes.Internal, "failed to load NodeInstanceGroupState (unscoped): %v", err)
	}

	return &model, nil
}

func (r *stateRepoImpl) LoadOrCreate(ctx context.Context, dataplaneName, nodeSpec, primaryInstance string) (*entity.DataplaneNodeInstanceGroupState, error) {
	if dataplaneName == "" || nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dataplaneName, nodeSpec, and primaryInstance are required")
	}
	nodeSpec = normalizeNodeSpec(nodeSpec)

	// Try to load existing (include soft-deleted records to handle Delete→LoadOrCreate sequences)
	existing, err := r.loadByDataplaneAndSpecUnscoped(ctx, dataplaneName, nodeSpec, primaryInstance)
	if err != nil {
		return nil, err
	}
	if existing != nil {
		if existing.DeletedAt.Valid {
			// Revive soft-deleted record
			wdb := r.provider.WriteDB(ctx)
			result := wdb.Unscoped().Model(&schema.DataplaneNodeInstanceGroupState{}).
				Where("id = ?", existing.ID).
				Updates(map[string]any{
					"deleted_at":       nil,
					"state":            schema.NodeInstanceGroupStateNormal,
					"state_entered_at": time.Now(),
					"fencing_token":    0,
					"cooldown_until":   nil,
				})
			if result.Error != nil {
				return nil, status.Errorf(codes.Internal, "failed to revive DataplaneNodeInstanceGroupState: %v", result.Error)
			}
			existing.DeletedAt.Valid = false
			existing.State = schema.NodeInstanceGroupStateNormal
			existing.NodeSpec = normalizeNodeSpecForOutput(existing.NodeSpec)
			return &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: *existing}, nil
		}
		existing.NodeSpec = normalizeNodeSpecForOutput(existing.NodeSpec)
		return &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: *existing}, nil
	}

	// Create new state in Normal
	wdb := r.provider.WriteDB(ctx)
	now := time.Now()
	model := &schema.DataplaneNodeInstanceGroupState{
		ID:              uuid.New().String(),
		DataplaneName:   dataplaneName,
		NodeSpec:        nodeSpec,
		PrimaryInstance: primaryInstance,
		State:           schema.NodeInstanceGroupStateNormal,
		StateEnteredAt:  now,
		FencingToken:    0,
	}

	if err := wdb.Create(model).Error; err != nil {
		// Handle race condition - another process might have created it
		if errors.Is(err, gorm.ErrDuplicatedKey) || gormutil.IsDuplicatedError(err) {
			return r.LoadByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, primaryInstance)
		}
		return nil, status.Errorf(codes.Internal, "failed to create DataplaneNodeInstanceGroupState: %v", err)
	}
	model.NodeSpec = normalizeNodeSpecForOutput(model.NodeSpec)

	return &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: *model}, nil
}

func (r *stateRepoImpl) ListByDataplane(ctx context.Context, dataplaneName string) ([]*entity.DataplaneNodeInstanceGroupState, error) {
	if dataplaneName == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dataplaneName is required")
	}

	rdb := r.provider.ReadDB(ctx)

	var models []schema.DataplaneNodeInstanceGroupState
	if err := rdb.Where("dataplane_name = ?", dataplaneName).Find(&models).Error; err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list NodeInstanceGroupStates: %v", err)
	}

	states := make([]*entity.DataplaneNodeInstanceGroupState, len(models))
	for i := range models {
		models[i].NodeSpec = normalizeNodeSpecForOutput(models[i].NodeSpec)
		states[i] = &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: models[i]}
	}

	return states, nil
}

func (r *stateRepoImpl) ListByState(ctx context.Context, state string) ([]*entity.DataplaneNodeInstanceGroupState, error) {
	rdb := r.provider.ReadDB(ctx)

	var models []schema.DataplaneNodeInstanceGroupState
	if err := rdb.Where("state = ?", state).Find(&models).Error; err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list NodeInstanceGroupStates by state: %v", err)
	}

	states := make([]*entity.DataplaneNodeInstanceGroupState, len(models))
	for i := range models {
		models[i].NodeSpec = normalizeNodeSpecForOutput(models[i].NodeSpec)
		states[i] = &entity.DataplaneNodeInstanceGroupState{DataplaneNodeInstanceGroupState: models[i]}
	}

	return states, nil
}

func (r *stateRepoImpl) Delete(ctx context.Context, id string) error {
	if id == "" {
		return status.Errorf(codes.InvalidArgument, "id is empty")
	}

	wdb := r.provider.WriteDB(ctx)

	if err := wdb.Where("id = ?", id).Delete(&schema.DataplaneNodeInstanceGroupState{}).Error; err != nil {
		return status.Errorf(codes.Internal, "failed to delete DataplaneNodeInstanceGroupState: %v", err)
	}

	return nil
}

func (r *stateRepoImpl) DeleteByDataplaneAndSpec(ctx context.Context, dataplaneName, nodeSpec, primaryInstance string) error {
	if dataplaneName == "" || nodeSpec == "" || primaryInstance == "" {
		return status.Errorf(codes.InvalidArgument, "dataplaneName, nodeSpec, and primaryInstance are required")
	}

	wdb := r.provider.WriteDB(ctx)
	normalized := normalizeNodeSpec(nodeSpec)

	if err := wdb.Where("dataplane_name = ? AND node_spec = ? AND primary_instance = ?", dataplaneName, normalized, primaryInstance).Delete(&schema.DataplaneNodeInstanceGroupState{}).Error; err != nil {
		return status.Errorf(codes.Internal, "failed to delete DataplaneNodeInstanceGroupState for %s/%s/%s: %v", dataplaneName, nodeSpec, primaryInstance, err)
	}

	return nil
}

func (r *stateRepoImpl) DeleteBySpecAndPrimary(ctx context.Context, nodeSpec, primaryInstance string) error {
	if nodeSpec == "" || primaryInstance == "" {
		return status.Errorf(codes.InvalidArgument, "nodeSpec and primaryInstance are required")
	}

	wdb := r.provider.WriteDB(ctx)
	normalized := normalizeNodeSpec(nodeSpec)

	if err := wdb.Where("node_spec = ? AND primary_instance = ?", normalized, primaryInstance).Delete(&schema.DataplaneNodeInstanceGroupState{}).Error; err != nil {
		return status.Errorf(codes.Internal, "failed to delete DataplaneNodeInstanceGroupStates by spec %s/%s: %v", nodeSpec, primaryInstance, err)
	}

	return nil
}

func (r *stateRepoImpl) TransitionState(ctx context.Context, id, expectedState, newState string) error {
	if id == "" {
		return status.Errorf(codes.InvalidArgument, "id is empty")
	}

	wdb := r.provider.WriteDB(ctx)
	now := time.Now()

	// Use optimistic locking via expected state check
	result := wdb.Model(&schema.DataplaneNodeInstanceGroupState{}).
		Where("id = ? AND state = ?", id, expectedState).
		Updates(map[string]any{
			"state":            newState,
			"state_entered_at": now,
			"fencing_token":    gorm.Expr("fencing_token + 1"),
			"cooldown_until":   nil, // Clear cooldown on successful transition
		})

	if result.Error != nil {
		return status.Errorf(codes.Internal, "failed to transition DataplaneNodeInstanceGroupState: %v", result.Error)
	}

	if result.RowsAffected == 0 {
		return status.Errorf(codes.FailedPrecondition, "NodeInstanceGroupState %s is not in expected state %s", id, expectedState)
	}

	return nil
}

func (r *stateRepoImpl) TransitionStateWithCooldown(ctx context.Context, id, expectedState, newState string, cooldownUntil *time.Time) error {
	if id == "" {
		return status.Errorf(codes.InvalidArgument, "id is empty")
	}

	wdb := r.provider.WriteDB(ctx)
	now := time.Now()

	result := wdb.Model(&schema.DataplaneNodeInstanceGroupState{}).
		Where("id = ? AND state = ?", id, expectedState).
		Updates(map[string]any{
			"state":            newState,
			"state_entered_at": now,
			"fencing_token":    gorm.Expr("fencing_token + 1"),
			"cooldown_until":   cooldownUntil,
		})

	if result.Error != nil {
		return status.Errorf(codes.Internal, "failed to transition DataplaneNodeInstanceGroupState: %v", result.Error)
	}

	if result.RowsAffected == 0 {
		return status.Errorf(codes.FailedPrecondition, "NodeInstanceGroupState %s is not in expected state %s", id, expectedState)
	}

	return nil
}

func (r *stateRepoImpl) SetCooldown(ctx context.Context, id string, until time.Time) error {
	if id == "" {
		return status.Errorf(codes.InvalidArgument, "id is empty")
	}

	wdb := r.provider.WriteDB(ctx)

	result := wdb.Model(&schema.DataplaneNodeInstanceGroupState{}).
		Where("id = ?", id).
		Update("cooldown_until", until)

	if result.Error != nil {
		return status.Errorf(codes.Internal, "failed to set cooldown: %v", result.Error)
	}

	if result.RowsAffected == 0 {
		return status.Errorf(codes.NotFound, "NodeInstanceGroupState not found: %s", id)
	}

	return nil
}

func (r *stateRepoImpl) ClearCooldown(ctx context.Context, id string) error {
	if id == "" {
		return status.Errorf(codes.InvalidArgument, "id is empty")
	}

	wdb := r.provider.WriteDB(ctx)

	result := wdb.Model(&schema.DataplaneNodeInstanceGroupState{}).
		Where("id = ?", id).
		Update("cooldown_until", nil)

	if result.Error != nil {
		return status.Errorf(codes.Internal, "failed to clear cooldown: %v", result.Error)
	}

	if result.RowsAffected == 0 {
		return status.Errorf(codes.NotFound, "NodeInstanceGroupState not found: %s", id)
	}

	return nil
}