package controlplane

import (
	"context"
	"encoding/json"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"aigw/internal/auth"
	"aigw/internal/catalog"
	"aigw/internal/domain"
)

const broadcastQueueSize = 128

type managerStore interface {
	LoadSnapshot(context.Context) (Snapshot, error)
	DatabaseGeneration(context.Context) (int64, error)
	PublishChange(context.Context, ChangeEvent) error
	Subscribe(context.Context) (<-chan ChangeMessage, func() error, error)
	RedisEnabled() bool
}

type policyReplacer interface {
	ReplacePolicies([]domain.LimitPolicy)
}

type Manager struct {
	store          managerStore
	catalog        *catalog.Catalog
	authenticator  *auth.StaticAuthenticator
	logger         *slog.Logger
	pollInterval   time.Duration
	generation     atomic.Int64
	redisConnected atomic.Bool
	reloadMu       sync.Mutex
	broadcasts     chan ChangeEvent
	policyTarget   policyReplacer
}

func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration, policyTargets ...policyReplacer) *Manager {
	if pollInterval <= 0 {
		pollInterval = 30 * time.Second
	}
	var policyTarget policyReplacer
	if len(policyTargets) > 0 {
		policyTarget = policyTargets[0]
	}
	return &Manager{
		store: store, catalog: modelCatalog, authenticator: authenticator,
		logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), policyTarget: policyTarget,
	}
}

func (m *Manager) Reload(ctx context.Context) (int64, error) {
	m.reloadMu.Lock()
	defer m.reloadMu.Unlock()
	snapshot, err := m.store.LoadSnapshot(ctx)
	if err != nil {
		return m.generation.Load(), err
	}
	m.catalog.Replace(snapshot.Models)
	m.authenticator.ReplaceHashed(snapshot.APIKeys)
	if m.policyTarget != nil {
		m.policyTarget.ReplacePolicies(snapshot.Limits)
	}
	m.generation.Store(snapshot.Generation)
	m.logger.Info("control_plane_reloaded", "generation", snapshot.Generation, "models", len(snapshot.Models), "api_keys", len(snapshot.APIKeys))
	return snapshot.Generation, nil
}

func (m *Manager) AfterMutation(ctx context.Context, generation int64, resource, id string) error {
	loadedGeneration, err := m.Reload(ctx)
	if err != nil {
		return err
	}
	if !m.store.RedisEnabled() {
		return nil
	}
	if loadedGeneration > generation {
		generation = loadedGeneration
	}
	event := newChange(generation, resource, id)
	select {
	case m.broadcasts <- event:
	default:
		m.logger.Warn("control_plane_publish_dropped", "generation", generation, "resource", resource, "id", id, "error", "broadcast queue is full")
	}
	return nil
}

func (m *Manager) Generation() int64 {
	return m.generation.Load()
}

func (m *Manager) RedisConfigured() bool {
	return m.store.RedisEnabled()
}

func (m *Manager) RedisConnected() bool {
	return m.redisConnected.Load()
}

func (m *Manager) Run(ctx context.Context) {
	var workers sync.WaitGroup
	if m.store.RedisEnabled() {
		workers.Add(2)
		go func() {
			defer workers.Done()
			m.runSubscriptions(ctx)
		}()
		go func() {
			defer workers.Done()
			m.runBroadcasts(ctx)
		}()
	} else {
		m.logger.Info("control_plane_redis_disabled", "fallback", "postgres_polling")
	}
	m.runPolling(ctx)
	workers.Wait()
}

func (m *Manager) runPolling(ctx context.Context) {
	ticker := time.NewTicker(m.pollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			generation, err := m.store.DatabaseGeneration(ctx)
			if err != nil {
				m.logger.Warn("control_plane_generation_check_failed", "error", err)
				continue
			}
			if generation > m.generation.Load() {
				if _, err := m.Reload(ctx); err != nil {
					m.logger.Error("control_plane_reload_failed", "source", "postgres", "error", err)
				}
			}
		}
	}
}

func (m *Manager) runSubscriptions(ctx context.Context) {
	retryDelay := m.pollInterval
	if retryDelay > time.Second {
		retryDelay = time.Second
	}
	if retryDelay <= 0 {
		retryDelay = time.Second
	}

	for ctx.Err() == nil {
		messages, closeSubscription, err := m.store.Subscribe(ctx)
		if err != nil {
			m.redisConnected.Store(false)
			m.logger.Warn("control_plane_subscription_failed", "error", err)
			if !waitForRetry(ctx, retryDelay) {
				return
			}
			continue
		}
		m.redisConnected.Store(true)
		m.logger.Info("control_plane_subscription_connected")

		closed := false
		for !closed {
			select {
			case <-ctx.Done():
				closed = true
			case message, ok := <-messages:
				if !ok {
					closed = true
					continue
				}
				var event ChangeEvent
				if json.Unmarshal([]byte(message.Payload), &event) != nil || event.Generation <= m.generation.Load() {
					continue
				}
				if _, err := m.Reload(ctx); err != nil {
					m.logger.Error("control_plane_reload_failed", "source", "redis", "error", err)
				}
			}
		}
		m.redisConnected.Store(false)
		if closeSubscription != nil {
			_ = closeSubscription()
		}
		if ctx.Err() == nil {
			m.logger.Warn("control_plane_subscription_disconnected")
			if !waitForRetry(ctx, retryDelay) {
				return
			}
		}
	}
}

func (m *Manager) runBroadcasts(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case event := <-m.broadcasts:
			publishContext, cancel := context.WithTimeout(ctx, 2*time.Second)
			err := m.store.PublishChange(publishContext, event)
			cancel()
			if err != nil {
				m.redisConnected.Store(false)
				m.logger.Warn("control_plane_publish_failed", "generation", event.Generation, "resource", event.Resource, "id", event.ID, "error", err)
			} else {
				m.redisConnected.Store(true)
			}
		}
	}
}

func waitForRetry(ctx context.Context, delay time.Duration) bool {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}