package encryptiondata

import (
	"encoding/base64"
	"fmt"
	"sort"
	"strings"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/runtime/schema"
	apiserverconfigv1 "k8s.io/apiserver/pkg/apis/apiserver/v1"
	"k8s.io/klog/v2"

	configv1 "github.com/openshift/api/config/v1"

	"github.com/openshift/library-go/pkg/operator/encryption/crypto"
	"github.com/openshift/library-go/pkg/operator/encryption/secrets"
	"github.com/openshift/library-go/pkg/operator/encryption/state"
)

var (
	emptyStaticKey = base64.StdEncoding.EncodeToString(crypto.NewEmptyKey())
)

// Config stores the EncryptionConfiguration and additional operator-managed
// encryption state that doesn't fit into the upstream type.
type Config struct {
	Encryption *apiserverconfigv1.EncryptionConfiguration
	// KMSPlugins maps keyID to plugin-specific configuration,
	// carried from Key Secrets into the encryption-config Secret.
	KMSPlugins map[string]configv1.KMSPluginConfig
	// KMSPluginsSecretData maps keyID to secret data carried from
	// Key Secrets into the encryption-config Secret.
	// Structure: keyID → referenceName → dataKey → value.
	KMSPluginsSecretData KMSPluginsReferenceData
	// KMSPluginsConfigMapData maps keyID to configmap data carried from
	// Key Secrets into the encryption-config Secret.
	// Structure: keyID → referenceName → dataKey → value.
	KMSPluginsConfigMapData KMSPluginsReferenceData
}

// KMSPluginsReferenceData maps keyID to reference data carried from Key Secrets into
// the encryption-config Secret. Structure: keyID → referenceName → dataKey → value.
type KMSPluginsReferenceData struct {
	byKeyID map[string]state.KMSReferenceData
}

// Get returns the KMSReferenceData for the given keyID and true if it exists,
// or a zero value and false otherwise. It is safe to call on a nil map.
func (d *KMSPluginsReferenceData) Get(keyID string) (state.KMSReferenceData, bool) {
	if d.byKeyID == nil {
		return state.KMSReferenceData{}, false
	}
	sd, ok := d.byKeyID[keyID]
	return sd, ok
}

// SetFromRawKey stores a value for the given keyID, splitting rawKey
// on "_" into referenceName and dataKey.
func (d *KMSPluginsReferenceData) SetFromRawKey(keyID, rawKey string, value []byte) error {
	if len(keyID) == 0 {
		return fmt.Errorf("keyID must not be empty")
	}
	if d.byKeyID == nil {
		d.byKeyID = map[string]state.KMSReferenceData{}
	}
	sd := d.byKeyID[keyID]
	if err := sd.SetFromRawKey(rawKey, value); err != nil {
		return err
	}
	d.byKeyID[keyID] = sd
	return nil
}

// FlatEntriesByKeyID returns the stored data as a map of keyID to flat entries,
// where each flat entry is keyed by "referenceName_dataKey".
func (d *KMSPluginsReferenceData) FlatEntriesByKeyID() map[string]map[string][]byte {
	if d.byKeyID == nil {
		return nil
	}
	result := map[string]map[string][]byte{}
	for keyID, sd := range d.byKeyID {
		if flat := sd.FlatEntries(); flat != nil {
			result[keyID] = flat
		}
	}
	return result
}

func (c *Config) HasEncryptionConfiguration() bool {
	return c != nil && c.Encryption != nil
}

// FromEncryptionState converts encryption state to Config.
func FromEncryptionState(encryptionState map[schema.GroupResource]state.GroupResourceState) (*Config, error) {
	resourceConfigs := make([]apiserverconfigv1.ResourceConfiguration, 0, len(encryptionState))
	var kmsPlugins map[string]configv1.KMSPluginConfig
	var kmsPluginsSecretData KMSPluginsReferenceData
	var kmsPluginsConfigMapData KMSPluginsReferenceData

	for gr, grKeys := range encryptionState {
		resourceConfigs = append(resourceConfigs, apiserverconfigv1.ResourceConfiguration{
			Resources: []string{gr.String()}, // we are forced to lose data here because this API is broken
			Providers: stateToProviders(gr.Resource, grKeys),
		})

		// Collect KMS plugin configs and secret data from read keys (which already
		// include the write key). We iterate over encryptionState which is keyed by
		// GroupResource, so the same keyID is seen once per resource. Since all
		// resources share the same Key Secret, the plugin config and secret data are
		// identical across duplicates and we only need to keep the first occurrence.
		for _, key := range grKeys.ReadKeys {
			if key.HasKMSPlugin() {
				if kmsPlugins == nil {
					kmsPlugins = map[string]configv1.KMSPluginConfig{}
				}
				if plugin, exists := kmsPlugins[key.Key.Name]; exists {
					// Sanity check: the same keyID seen from a different resource must carry
					// an identical plugin config, since they originate from the same Key Secret.
					if !equality.Semantic.DeepEqual(plugin, key.KMS.Plugin) {
						return nil, fmt.Errorf("KMS plugin config mismatch for keyID %s: configs from different resources must be identical", key.Key.Name)
					}
				} else {
					kmsPlugins[key.Key.Name] = key.KMS.Plugin
				}
			}
			if key.HasKMSSecretData() {
				if existing, exists := kmsPluginsSecretData.Get(key.Key.Name); exists {
					if !equality.Semantic.DeepEqual(existing.FlatEntries(), key.KMS.PluginSecretData.FlatEntries()) {
						return nil, fmt.Errorf("KMS secret data mismatch for keyID %s: secret data from different resources must be identical", key.Key.Name)
					}
				} else {
					for rawKey, value := range key.KMS.PluginSecretData.FlatEntries() {
						if err := kmsPluginsSecretData.SetFromRawKey(key.Key.Name, rawKey, value); err != nil {
							return nil, fmt.Errorf("failed to copy secret data for keyID %s: %w", key.Key.Name, err)
						}
					}
				}
			}
			if key.HasKMSConfigMapData() {
				if existing, exists := kmsPluginsConfigMapData.Get(key.Key.Name); exists {
					if !equality.Semantic.DeepEqual(existing.FlatEntries(), key.KMS.PluginConfigMapData.FlatEntries()) {
						return nil, fmt.Errorf("KMS configmap data mismatch for keyID %s: configmap data from different resources must be identical", key.Key.Name)
					}
				} else {
					for rawKey, value := range key.KMS.PluginConfigMapData.FlatEntries() {
						if err := kmsPluginsConfigMapData.SetFromRawKey(key.Key.Name, rawKey, value); err != nil {
							return nil, fmt.Errorf("failed to copy configmap data for keyID %s: %w", key.Key.Name, err)
						}
					}
				}
			}
		}
	}

	// make sure our output is stable
	sort.Slice(resourceConfigs, func(i, j int) bool {
		return resourceConfigs[i].Resources[0] < resourceConfigs[j].Resources[0] // each resource has its own keys
	})

	return &Config{
		Encryption:              &apiserverconfigv1.EncryptionConfiguration{Resources: resourceConfigs},
		KMSPlugins:              kmsPlugins,
		KMSPluginsSecretData:    kmsPluginsSecretData,
		KMSPluginsConfigMapData: kmsPluginsConfigMapData,
	}, nil
}

// ToEncryptionState converts config to state.
// Read keys contain a potential write key. Read keys are sorted, recent first.
//
// It assumes:
// - the first provider provides the write key
// - the structure of the encryptionConfig matches the output generated by FromEncryptionState:
//   - one resource per provider
//   - one key per provider
//   - each resource has a distinct configuration with zero or more key based providers and the identity provider.
//   - the last providers might be of type aesgcm. Then it carries the names of identity keys, recent first.
//     We never use aesgcm as a real key because it is unsafe.
func ToEncryptionState(secretData *Config, keySecrets []*corev1.Secret) (map[schema.GroupResource]state.GroupResourceState, []state.KeyState) {
	backedKeys := make([]state.KeyState, 0, len(keySecrets))
	for _, s := range keySecrets {
		km, err := secrets.ToKeyState(s)
		if err != nil {
			klog.Warningf("skipping invalid secret: %v", err)
			continue
		}
		km.Backed = true
		backedKeys = append(backedKeys, km)
	}
	backedKeys = state.SortRecentFirst(backedKeys)

	if !secretData.HasEncryptionConfiguration() {
		return nil, backedKeys
	}

	out := map[schema.GroupResource]state.GroupResourceState{}
	for _, resourceConfig := range secretData.Encryption.Resources {
		// resources should be a single group resource
		if len(resourceConfig.Resources) != 1 {
			klog.Warningf("skipping invalid encryption config for resource %s", resourceConfig.Resources)
			continue // should never happen
		}

		grState := state.GroupResourceState{}

		for i, provider := range resourceConfig.Providers {
			var ks state.KeyState

			switch {
			case provider.AESCBC != nil && len(provider.AESCBC.Keys) == 1:
				ks = state.KeyState{
					Key:  provider.AESCBC.Keys[0],
					Mode: state.AESCBC,
				}

			case provider.Secretbox != nil && len(provider.Secretbox.Keys) == 1:
				ks = state.KeyState{
					Key:  provider.Secretbox.Keys[0],
					Mode: state.SecretBox,
				}

			case provider.Identity != nil:
				// skip fake provider. If this is write-key, wait for first aesgcm provider providing the write key.
				continue

			case provider.AESGCM != nil && len(provider.AESGCM.Keys) == 1:
				s := state.AESGCM
				if provider.AESGCM.Keys[0].Secret == emptyStaticKey {
					s = state.Identity
				}

				ks = state.KeyState{
					Key:  provider.AESGCM.Keys[0],
					Mode: s,
				}

			case provider.KMS != nil:
				// Name and Secret must match to find our backed Secret
				keyID, err := getKeyIDFromPluginName(provider.KMS.Name)
				if err != nil {
					klog.Warningf("resource: %s provider index: %d: Skipping invalid plugin name %s: %v", resourceConfig.Resources[0], i, provider.KMS.Name, err)
					continue // should never happen
				}
				ks = state.KeyState{
					Key: apiserverconfigv1.Key{
						Name:   keyID,
						Secret: emptyStaticKey,
					},
					Mode: state.KMS,
				}
			default:
				klog.Infof("skipping invalid provider index %d for resource %s", i, resourceConfig.Resources[0])
				continue // should never happen
			}

			// enrich KeyState with values from secrets
			for _, k := range backedKeys {
				if state.EqualKeyAndEqualID(&ks, &k) {
					ks = k
					break
				}
			}

			if i == 0 || (ks.Mode == state.Identity && !grState.HasWriteKey()) {
				grState.WriteKey = ks
			}

			grState.ReadKeys = append(grState.ReadKeys, ks) // also for write key as they are also read keys
		}

		// sort read-keys, recent first
		grState.ReadKeys = state.SortRecentFirst(grState.ReadKeys)

		out[schema.ParseGroupResource(resourceConfig.Resources[0])] = grState
	}

	return out, backedKeys
}

// stateToProviders maps the write and read secrets to the equivalent read and write keys.
// it primarily handles the conversion of KeyState to the appropriate provider config.
// the identity mode is transformed into a custom aesgcm provider that simply exists to
// curry the associated null key secret through the encryption state machine.
func stateToProviders(resource string, desired state.GroupResourceState) []apiserverconfigv1.ProviderConfiguration {
	allKeys := desired.ReadKeys

	providers := make([]apiserverconfigv1.ProviderConfiguration, 0, len(allKeys)+1) // one extra for identity

	// Write key comes first. Filter it out in the tail of read keys.
	if desired.HasWriteKey() {
		allKeys = append([]state.KeyState{desired.WriteKey}, allKeys...)
		for i := 1; i < len(allKeys); i++ {
			if state.EqualKeyAndEqualID(&allKeys[i], &desired.WriteKey) {
				allKeys = append(allKeys[:i], allKeys[i+1:]...)
				break
			}
		}
	} else {
		// no write key => identity write key
		providers = append(providers, apiserverconfigv1.ProviderConfiguration{
			Identity: &apiserverconfigv1.IdentityConfiguration{},
		})
	}

	aesgcmProviders := []apiserverconfigv1.ProviderConfiguration{}
	for i, key := range allKeys {
		switch key.Mode {
		case state.AESCBC:
			providers = append(providers, apiserverconfigv1.ProviderConfiguration{
				AESCBC: &apiserverconfigv1.AESConfiguration{
					Keys: []apiserverconfigv1.Key{key.Key},
				},
			})
		case state.AESGCM:
			providers = append(providers, apiserverconfigv1.ProviderConfiguration{
				AESGCM: &apiserverconfigv1.AESConfiguration{
					Keys: []apiserverconfigv1.Key{key.Key},
				},
			})
		case state.SecretBox:
			providers = append(providers, apiserverconfigv1.ProviderConfiguration{
				Secretbox: &apiserverconfigv1.SecretboxConfiguration{
					Keys: []apiserverconfigv1.Key{key.Key},
				},
			})
		case state.Identity:
			if i == 0 {
				providers = append(providers, apiserverconfigv1.ProviderConfiguration{
					Identity: &apiserverconfigv1.IdentityConfiguration{},
				})
			}
			aesgcmProviders = append(aesgcmProviders, apiserverconfigv1.ProviderConfiguration{
				AESGCM: &apiserverconfigv1.AESConfiguration{
					Keys: []apiserverconfigv1.Key{key.Key},
				},
			})
		case state.KMS:
			if !key.HasKMSEncryption() {
				klog.Infof("skipping key %s for %s in KMS mode as its either KMS or KMS.Encryption is nil", key.Key.Name, resource)
				continue // this should never happen
			}
			// In order to preserve the uniqueness, we should insert resource name
			kmsCopy := key.KMS.Encryption.DeepCopy()
			kmsCopy.Name = createKMSPluginName(key.Key.Name, resource)
			provider := apiserverconfigv1.ProviderConfiguration{
				KMS: kmsCopy,
			}
			providers = append(providers, provider)
		default:
			// this should never happen because our input should always be valid
			klog.Infof("skipping key %s as it has invalid mode %s", key.Key.Name, key.Mode)
		}
	}

	// add fallback identity provider.
	if providers[0].Identity == nil {
		providers = append(providers, apiserverconfigv1.ProviderConfiguration{
			Identity: &apiserverconfigv1.IdentityConfiguration{},
		})
	}

	// add fake aesgm providers carrying identity names
	providers = append(providers, aesgcmProviders...)

	return providers
}

func createKMSPluginName(keyID, resource string) string {
	// Ideally we should have used keyID simply in kms plugin name.
	// However, this is an upstream constraint that every provider name must be unique.
	// To maintain uniqueness while still allowing access to the keyID, we generate plugin name in this format.
	return fmt.Sprintf("%s_%s", keyID, resource)
}

func getKeyIDFromPluginName(pluginName string) (string, error) {
	// We just need to obtain the keyID to find our backed secret
	// e.g. "1_secrets"
	parsed := strings.SplitN(pluginName, "_", 2)
	if len(parsed) != 2 {
		return "", fmt.Errorf("invalid plugin name %q: expected format keyID_resourceName", pluginName)
	}
	return parsed[0], nil
}
