package store

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/tinode/chat/server/logs"
	"github.com/tinode/chat/server/store/types"
)

type cachedMessageKey struct {
	enabled bool
	version int
	status  string
	key     []byte
	expires time.Time
}

const messageKeyCacheTTL = 5 * time.Minute

var messageKeyCache = struct {
	sync.RWMutex
	items map[string]cachedMessageKey
}{items: make(map[string]cachedMessageKey)}

func messageKeyCacheID(tenantID types.TenantID, version int) string {
	return strconv.FormatInt(int64(tenantID), 10) + ":" + strconv.Itoa(version)
}

// WarmMessageEncryptionKeys preloads active keys. Missing keys leave the tenant in plaintext mode.
func WarmMessageEncryptionKeys() {
	if adminInternalConfigProvider == nil || Tenants == nil {
		return
	}
	tenants, err := Tenants.ListActive(10000)
	if err != nil {
		logs.Warn.Printf("message crypto: tenant preload failed: %v", err)
		return
	}
	for i := range tenants {
		if _, err = loadMessageKey(&tenants[i], 0); err != nil {
			logs.Info.Printf("message crypto: plaintext mode tenant=%s reason=%v", tenants[i].Code, err)
		}
	}
}

func loadMessageKey(tenant *types.Tenant, version int) (cachedMessageKey, error) {
	var zero cachedMessageKey
	if tenant == nil || adminInternalConfigProvider == nil {
		return zero, types.ErrNotFound
	}
	result, err := adminInternalConfigProvider.MessageKey(tenant.Code, version)
	if err != nil || result == nil {
		if err == nil {
			err = types.ErrNotFound
		}
		return zero, err
	}
	decoded, err := base64.StdEncoding.DecodeString(result.Key)
	if err != nil || len(decoded) != 32 {
		return zero, errors.New("invalid tenant message key")
	}
	entry := cachedMessageKey{enabled: result.Enabled, version: result.Version, status: result.Status, key: decoded, expires: time.Now().Add(messageKeyCacheTTL)}
	messageKeyCache.Lock()
	messageKeyCache.items[messageKeyCacheID(tenant.ID, result.Version)] = entry
	if version == 0 {
		messageKeyCache.items[messageKeyCacheID(tenant.ID, 0)] = entry
	}
	messageKeyCache.Unlock()
	return entry, nil
}

func getMessageKey(tenantID types.TenantID, version int) (cachedMessageKey, error) {
	cacheID := messageKeyCacheID(tenantID, version)
	messageKeyCache.RLock()
	entry, ok := messageKeyCache.items[cacheID]
	messageKeyCache.RUnlock()
	if ok && time.Now().Before(entry.expires) {
		return entry, nil
	}
	if ok {
		messageKeyCache.Lock()
		delete(messageKeyCache.items, cacheID)
		messageKeyCache.Unlock()
	}
	tenant, err := Tenants.GetByID(tenantID)
	if err != nil {
		return cachedMessageKey{}, err
	}
	entry, err = loadMessageKey(tenant, version)
	if err != nil && version == 0 {
		entry = cachedMessageKey{expires: time.Now().Add(messageKeyCacheTTL)}
		messageKeyCache.Lock()
		messageKeyCache.items[cacheID] = entry
		messageKeyCache.Unlock()
		return entry, nil
	}
	return entry, err
}

// PrepareMessageEncryptionKey verifies a pending key in memory and acknowledges
// it to Admin Server. Admin activates the key only after this succeeds.
func PrepareMessageEncryptionKey(tenantID types.TenantID, tenantCode string, version int) error {
	if version <= 0 || adminInternalConfigProvider == nil {
		return types.ErrMalformed
	}
	var tenant *types.Tenant
	var err error
	if tenantID.IsValid() {
		tenant, err = Tenants.GetByID(tenantID)
	} else {
		tenant, err = Tenants.GetByCode(strings.TrimSpace(tenantCode))
	}
	if err != nil || tenant == nil {
		return err
	}
	entry, err := loadMessageKey(tenant, version)
	if err != nil {
		return err
	}
	if strings.EqualFold(entry.status, "pending") {
		if err = adminInternalConfigProvider.MessageKeyAck(tenant.Code, version); err != nil {
			return err
		}
	}
	messageKeyCache.Lock()
	messageKeyCache.items[messageKeyCacheID(tenant.ID, 0)] = entry
	messageKeyCache.Unlock()
	return nil
}

// InvalidateMessageEncryptionKeys forces the next access to reload tenant keys.
func InvalidateMessageEncryptionKeys(tenantID types.TenantID) {
	prefix := strconv.FormatInt(int64(tenantID), 10) + ":"
	messageKeyCache.Lock()
	for key := range messageKeyCache.items {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			delete(messageKeyCache.items, key)
		}
	}
	messageKeyCache.Unlock()
}

func messageAAD(msg *types.Message, keyVersion int) []byte {
	topic := []byte(msg.Topic)
	aad := make([]byte, 8+8+4+len(topic))
	binary.BigEndian.PutUint64(aad[0:8], uint64(msg.TenantID))
	binary.BigEndian.PutUint64(aad[8:16], uint64(msg.SeqId))
	binary.BigEndian.PutUint32(aad[16:20], uint32(keyVersion))
	copy(aad[20:], topic)
	return aad
}

func encryptMessageContent(msg *types.Message) (*types.Message, error) {
	entry, err := getMessageKey(msg.TenantID, 0)
	if err != nil || !entry.enabled || entry.version <= 0 {
		if err != nil {
			logs.Warn.Printf("message crypto: active key unavailable, storing plaintext tenant=%d reason=%v", msg.TenantID, err)
		}
		copy := *msg
		copy.KeyVersion, copy.ContentNonce = 0, nil
		return &copy, nil
	}
	plain, err := json.Marshal(msg.Content)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(entry.key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	copy := *msg
	copy.KeyVersion = entry.version
	copy.ContentNonce = nonce
	copy.Content = base64.StdEncoding.EncodeToString(gcm.Seal(nil, nonce, plain, messageAAD(&copy, entry.version)))
	return &copy, nil
}

func decryptMessageContent(msg *types.Message) error {
	// 撤回消息的内容已被清空（recall 会把 head/content 置 NULL，但 content_nonce/key_version 残留），
	// 此类行没有密文可解，直接放行让上层按已撤回形态下发，否则整个 get data 窗口会被单条撤回消息毒化成 500。
	if msg == nil || msg.KeyVersion <= 0 || msg.Content == nil {
		return nil
	}
	entry, err := getMessageKey(msg.TenantID, msg.KeyVersion)
	if err != nil {
		return fmt.Errorf("message key version %d unavailable: %w", msg.KeyVersion, err)
	}
	ciphertextText, ok := msg.Content.(string)
	if !ok {
		return errors.New("encrypted message content is not a string")
	}
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextText)
	if err != nil {
		return err
	}
	block, err := aes.NewCipher(entry.key)
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	if len(msg.ContentNonce) != gcm.NonceSize() {
		return fmt.Errorf("invalid message nonce length: got %d want %d", len(msg.ContentNonce), gcm.NonceSize())
	}
	plain, err := gcm.Open(nil, msg.ContentNonce, ciphertext, messageAAD(msg, msg.KeyVersion))
	if err != nil {
		return err
	}
	return json.Unmarshal(plain, &msg.Content)
}

func decryptMessages(messages []types.Message) error {
	for i := range messages {
		if err := decryptMessageContent(&messages[i]); err != nil {
			return err
		}
	}
	return nil
}

func messageContentMatchesSearch(content any, query string) bool {
	query = strings.ToLower(strings.TrimSpace(query))
	if query == "" {
		return true
	}
	raw, err := json.Marshal(content)
	if err != nil {
		return false
	}
	text := strings.ToLower(string(raw))
	if strings.Contains(text, query) {
		return true
	}
	for _, term := range strings.Fields(query) {
		if !strings.Contains(text, term) {
			return false
		}
	}
	return true
}