package discovery

import (
	"encoding/json"
	"errors"
	"strings"

	"github.com/nats-io/nats.go"
	"github.com/tech/sendico/pkg/mlogger"
	"go.uber.org/zap"
)

const DefaultKVBucket = "discovery_registry"

type KVStore struct {
	logger mlogger.Logger
	kv     nats.KeyValue
	bucket string
}

func NewKVStore(logger mlogger.Logger, js nats.JetStreamContext, bucket string) (*KVStore, error) {
	if js == nil {
		return nil, errors.New("discovery kv: jetstream is nil")
	}
	if logger != nil {
		logger = logger.Named("discovery_kv")
	}
	bucket = strings.TrimSpace(bucket)
	if bucket == "" {
		bucket = DefaultKVBucket
	}
	kv, err := js.KeyValue(bucket)
	if err != nil {
		if errors.Is(err, nats.ErrBucketNotFound) {
			kv, err = js.CreateKeyValue(&nats.KeyValueConfig{
				Bucket:      bucket,
				Description: "service discovery registry",
				History:     1,
			})
			if err == nil && logger != nil {
				logger.Info("Discovery KV bucket created", zap.String("bucket", bucket))
			}
		}
		if err != nil {
			return nil, err
		}
	}

	return &KVStore{
		logger: logger,
		kv:     kv,
		bucket: bucket,
	}, nil
}

func (s *KVStore) Put(entry RegistryEntry) error {
	if s == nil || s.kv == nil {
		return errors.New("discovery kv: not configured")
	}
	key := registryEntryKey(normalizeEntry(entry))
	if key == "" {
		return errors.New("discovery kv: entry key is empty")
	}
	payload, err := json.Marshal(entry)
	if err != nil {
		return err
	}
	_, err = s.kv.Put(kvKeyFromRegistryKey(key), payload)
	if err != nil && s.logger != nil {
		fields := append(entryFields(entry), zap.String("bucket", s.bucket), zap.String("key", key), zap.Error(err))
		s.logger.Warn("Failed to persist discovery entry", fields...)
	}
	return err
}

func (s *KVStore) Delete(id string) error {
	if s == nil || s.kv == nil {
		return errors.New("discovery kv: not configured")
	}
	key := kvKeyFromRegistryKey(id)
	if key == "" {
		return nil
	}
	if err := s.kv.Delete(key); err != nil && s.logger != nil {
		s.logger.Warn("Failed to delete discovery entry", zap.String("bucket", s.bucket), zap.String("key", key), zap.Error(err))
		return err
	}
	return nil
}

func (s *KVStore) WatchAll() (nats.KeyWatcher, error) {
	if s == nil || s.kv == nil {
		return nil, errors.New("discovery kv: not configured")
	}
	return s.kv.WatchAll()
}

func (s *KVStore) Bucket() string {
	if s == nil {
		return ""
	}
	return s.bucket
}