package app

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/store/app/registry"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

const CredentialStorageScheme = model.CredentialStorageScheme

func HasCredentialStorage(key []byte) bool {
	return len(key) == 32
}

func MarshalCredentialsObject(key []byte, value map[string]any) ([]byte, error) {
	encoded, err := registry.MarshalJSONObject(value, "credentials")
	if err != nil {
		return nil, err
	}
	return EncryptStoredCredentialsJSON(key, encoded)
}

func EncryptStoredCredentialsJSON(key []byte, plaintextJSON []byte) ([]byte, error) {
	normalized, err := NormalizeStoredCredentialsJSON(plaintextJSON)
	if err != nil {
		return nil, err
	}
	if JSONObjectIsEmpty(normalized) {
		return []byte("{}"), nil
	}
	if !HasCredentialStorage(key) {
		return nil, badRequest("app credential storage is not configured on server")
	}

	ciphertext, err := EncryptCredential(key, string(normalized))
	if err != nil {
		return nil, fmt.Errorf("encrypt app credentials: %w", err)
	}
	envelope := model.CredentialEnvelope{
		Storage:    CredentialStorageScheme,
		Ciphertext: ciphertext,
	}
	encodedEnvelope, err := model.MarshalJSONObject("encrypted app credential envelope", envelope)
	if err != nil {
		return nil, fmt.Errorf("encode encrypted app credential envelope: %w", err)
	}
	return encodedEnvelope, nil
}

func DecryptStoredCredentialsJSON(key []byte, storedJSON []byte) ([]byte, error) {
	return DecryptStoredCredentialsJSONWithFallback(key, nil, storedJSON)
}

func DecryptStoredCredentialsJSONWithFallback(key []byte, legacyKeys [][]byte, storedJSON []byte) ([]byte, error) {
	normalized, err := NormalizeStoredCredentialsJSON(storedJSON)
	if err != nil {
		return nil, err
	}

	var envelope model.CredentialEnvelope
	if err := model.UnmarshalJSONObject("encrypted app credential envelope", normalized, &envelope); err != nil {
		return nil, fmt.Errorf("decode encrypted app credential envelope: %w", err)
	}
	if strings.TrimSpace(envelope.Storage) == "" {
		return normalized, nil
	}
	if strings.TrimSpace(envelope.Storage) != CredentialStorageScheme {
		return nil, fmt.Errorf("unsupported app credential storage scheme %q", envelope.Storage)
	}
	keys := credentialStorageKeyCandidates(key, legacyKeys)
	if len(keys) == 0 {
		return nil, badRequest("app credentials are configured but server credential storage is unavailable")
	}

	plaintext, err := DecryptCredentialWithFallback(keys, envelope.Ciphertext)
	if err != nil {
		return nil, fmt.Errorf("decrypt app credentials: %w", err)
	}
	return NormalizeStoredCredentialsJSON([]byte(plaintext))
}

func NormalizeStoredCredentialsJSON(payload []byte) ([]byte, error) {
	trimmedPayload := strings.TrimSpace(string(payload))
	if trimmedPayload == "" {
		return []byte("{}"), nil
	}

	var decoded map[string]any
	if err := model.UnmarshalJSONObject("app credentials", []byte(trimmedPayload), &decoded); err != nil {
		return nil, fmt.Errorf("decode app credentials: %w", err)
	}
	if decoded == nil {
		return []byte("{}"), nil
	}

	normalized, err := model.MarshalJSONObject("app credentials", decoded)
	if err != nil {
		return nil, fmt.Errorf("encode app credentials: %w", err)
	}
	return normalized, nil
}

func JSONObjectIsEmpty(payload []byte) bool {
	var decoded map[string]any
	if err := model.UnmarshalJSONObject("app credentials", payload, &decoded); err != nil {
		return false
	}
	return len(decoded) == 0
}

func EncryptCredential(key []byte, plaintext string) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("create app credential cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("create app credential gcm: %w", err)
	}

	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", fmt.Errorf("generate app credential nonce: %w", err)
	}

	sealed := aead.Seal(nonce, nonce, []byte(plaintext), nil)
	return base64.RawStdEncoding.EncodeToString(sealed), nil
}

func DecryptCredential(key []byte, ciphertext string) (string, error) {
	return DecryptCredentialWithFallback([][]byte{key}, ciphertext)
}

func DecryptCredentialWithFallback(keys [][]byte, ciphertext string) (string, error) {
	var firstErr error
	for _, key := range keys {
		if !HasCredentialStorage(key) {
			if firstErr == nil {
				firstErr = badRequest("app credentials are configured but server credential storage is unavailable")
			}
			continue
		}
		plaintext, err := decryptCredentialWithKey(key, ciphertext)
		if err == nil {
			return plaintext, nil
		}
		if firstErr == nil {
			firstErr = err
		}
	}
	if firstErr != nil {
		return "", firstErr
	}
	return "", badRequest("app credentials are configured but server credential storage is unavailable")
}

func decryptCredentialWithKey(key []byte, ciphertext string) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("create app credential cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("create app credential gcm: %w", err)
	}

	payload, err := base64.RawStdEncoding.DecodeString(strings.TrimSpace(ciphertext))
	if err != nil {
		return "", fmt.Errorf("decode app credential ciphertext: %w", err)
	}
	if len(payload) < aead.NonceSize() {
		return "", fmt.Errorf("app credential ciphertext is too short")
	}

	nonce := payload[:aead.NonceSize()]
	encrypted := payload[aead.NonceSize():]
	plaintext, err := aead.Open(nil, nonce, encrypted, nil)
	if err != nil {
		return "", fmt.Errorf("decrypt app credential ciphertext: %w", err)
	}
	return string(plaintext), nil
}