package discovery

import (
	"encoding/json"
	"sync"
	"time"

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

type RegistryWatcher struct {
	logger   mlogger.Logger
	registry *Registry
	kv       *KVStore
	watcher  nats.KeyWatcher

	stopOnce sync.Once
}

func NewRegistryWatcher(logger mlogger.Logger, msgBroker mb.Broker, registry *Registry) (*RegistryWatcher, error) {
	if msgBroker == nil {
		return nil, merrors.InvalidArgument("discovery watcher: broker is nil")
	}
	if registry == nil {
		registry = NewRegistry()
	}
	if logger == nil {
		return nil, merrors.InvalidArgument("discovery logger: logger must be provided")
	}
	logger = logger.Named("discovery_watcher")
	provider, ok := msgBroker.(jetStreamProvider)
	if !ok {
		return nil, merrors.Internal("discovery watcher: jetstream not available")
	}
	js := provider.JetStream()
	if js == nil {
		return nil, merrors.Internal("discovery watcher: jetstream not configured")
	}
	store, err := NewKVStore(logger, js, "")
	if err != nil {
		return nil, err
	}

	return &RegistryWatcher{
		logger:   logger,
		registry: registry,
		kv:       store,
	}, nil
}

func (w *RegistryWatcher) Start() error {
	if w == nil || w.kv == nil {
		return merrors.Internal("discovery watcher: not configured")
	}
	watcher, err := w.kv.WatchAll()
	if err != nil {
		return err
	}
	w.watcher = watcher
	w.logger.Info("Discovery registry watcher started", zap.String("bucket", w.kv.Bucket()))
	go w.consume(watcher)
	return nil
}

func (w *RegistryWatcher) Stop() {
	if w == nil {
		return
	}
	w.stopOnce.Do(func() {
		if w.watcher != nil {
			_ = w.watcher.Stop()
		}
		w.logger.Info("Discovery registry watcher stopped")
	})
}

func (w *RegistryWatcher) Registry() *Registry {
	if w == nil {
		return nil
	}
	return w.registry
}

func (w *RegistryWatcher) consume(watcher nats.KeyWatcher) {
	if w == nil || watcher == nil {
		return
	}
	initial := true
	initialCount := 0
	for entry := range watcher.Updates() {
		if entry == nil {
			if initial {
				fields := []zap.Field{zap.Int("entries", initialCount)}
				if w.kv != nil {
					fields = append(fields, zap.String("bucket", w.kv.Bucket()))
				}
				w.logger.Info("Discovery registry watcher initial sync complete", fields...)
				initial = false
			}
			continue
		}
		if initial && entry.Operation() == nats.KeyValuePut {
			initialCount++
		}
		switch entry.Operation() {
		case nats.KeyValueDelete, nats.KeyValuePurge:
			key := registryKeyFromKVKey(entry.Key())
			if key != "" {
				if w.registry.Delete(key) {
					w.logger.Info("Discovery registry entry removed", zap.String("key", key))
				}
			}
			continue
		case nats.KeyValuePut:
		default:
			continue
		}

		var payload RegistryEntry
		if err := json.Unmarshal(entry.Value(), &payload); err != nil {
			w.logger.Warn("Failed to decode discovery KV entry", zap.String("key", entry.Key()), zap.Error(err))
			continue
		}
		result := w.registry.UpsertEntry(payload, time.Now())
		if result.IsNew || result.BecameHealthy {
			fields := append(entryFields(result.Entry), zap.Bool("is_new", result.IsNew), zap.Bool("became_healthy", result.BecameHealthy))
			w.logger.Info("Discovery registry entry updated from KV", fields...)
		}
	}
}