package br

import (
	"context"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	encryptionkeyrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/encryptionkey"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

// Snapshot is the backup-time encryption snapshot persisted on backup records.
type Snapshot struct {
	CustomerManaged bool
	EncryptionKeyId uint64
}

// Resolver resolves backup encryption snapshots from cluster and encryption key metadata.
type Resolver struct {
	encryptionKeyRepo encryptionkeyrepo.Repo
}

// NewResolver creates a Resolver backed by encryption key repository.
func NewResolver(encryptionKeyRepo encryptionkeyrepo.Repo) *Resolver {
	return &Resolver{
		encryptionKeyRepo: encryptionKeyRepo,
	}
}

// ResolveByCluster resolves the encryption snapshot for a single cluster.
func (r *Resolver) ResolveByCluster(ctx context.Context, cluster *entity.Cluster) (Snapshot, bizerr.BizError) {
	if cluster == nil || cluster.EncryptionKeyId == 0 || r == nil || r.encryptionKeyRepo == nil {
		return Snapshot{}, nil
	}

	keyMap, berr := r.loadEncryptionKeys(ctx, []uint64{cluster.EncryptionKeyId})
	if berr != nil {
		return Snapshot{}, berr
	}
	return buildSnapshot(ctx, cluster, keyMap[cluster.EncryptionKeyId]), nil
}

// ResolveByClusters resolves encryption snapshots for the given clusters.
func (r *Resolver) ResolveByClusters(ctx context.Context, clusters map[uint64]*entity.Cluster) (map[uint64]Snapshot, bizerr.BizError) {
	snapshots := make(map[uint64]Snapshot, len(clusters))
	if len(clusters) == 0 || r == nil || r.encryptionKeyRepo == nil {
		return snapshots, nil
	}

	keyIds := make([]uint64, 0, len(clusters))
	for _, cluster := range clusters {
		if cluster.EncryptionKeyId == 0 {
			continue
		}
		keyIds = append(keyIds, cluster.EncryptionKeyId)
	}

	keyMap, berr := r.loadEncryptionKeys(ctx, keyIds)
	if berr != nil {
		return nil, berr
	}

	for clusterId, cluster := range clusters {
		if cluster == nil {
			tilog.Warn(ctx, "cluster not found while resolving backup encryption snapshot", zap.Uint64("cluster_id", clusterId))
			continue
		}
		snapshots[clusterId] = buildSnapshot(ctx, cluster, keyMap[cluster.EncryptionKeyId])
	}
	return snapshots, nil
}

func (r *Resolver) loadEncryptionKeys(ctx context.Context, keyIds []uint64) (map[uint64]*entity.EncryptionKey, bizerr.BizError) {
	keyMap := make(map[uint64]*entity.EncryptionKey, len(keyIds))
	if len(keyIds) == 0 || r == nil || r.encryptionKeyRepo == nil {
		return keyMap, nil
	}

	keys, berr := r.encryptionKeyRepo.ListEncryptionKeys(ctx, &encryptionkeyrepo.Filter{Ids: keyIds})
	if berr != nil {
		return nil, berr
	}
	for _, key := range keys {
		keyMap[key.Id] = key
	}
	return keyMap, nil
}

func buildSnapshot(ctx context.Context, cluster *entity.Cluster, key *entity.EncryptionKey) Snapshot {
	if cluster == nil || cluster.EncryptionKeyId == 0 {
		return Snapshot{}
	}

	snapshot := Snapshot{
		CustomerManaged: false,
		EncryptionKeyId: cluster.EncryptionKeyId,
	}
	if key == nil {
		tilog.Warn(ctx, "fall back to smek because encryption key details are missing",
			zap.Uint64("cluster_id", cluster.GetClusterId()),
			zap.Uint64("encryption_key_id", cluster.EncryptionKeyId))
		return snapshot
	}
	if key.CustomerManaged {
		snapshot.CustomerManaged = true
	}
	return snapshot
}