package federationapi

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"slices"
	"strconv"
	"sync"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/matrix/id"
	"github.com/dariuszjaniak/warp/server/internal/roomserver"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// The outbound queue: every event and EDU this server owes another one.
//
// It has three moving parts. The roomserver hands it every persisted event
// through EventPersisted, which only appends to an in-memory list — it is
// called with the room's write lock held and must not touch the database.
// A dispatcher drains that list, works out each event's destinations and
// writes one queue row per destination. And one worker per destination
// with something queued sends it, in transactions of up to 50 PDUs and
// 100 EDUs, backing off exponentially when the destination fails.
//
// The rows are deleted only once the destination has answered, so a crash
// or a restart loses nothing that had reached the database, and a
// destination that is down gets everything it missed when it comes back.
// What it does lose is the dispatcher's in-memory list: the events
// persisted in the instant before a crash. A remote server notices those
// as gaps the next time it hears from this one, and fetches them.

// Transaction limits, from the spec.
const (
	maxPDUsPerTxn = 50
	maxEDUsPerTxn = 100
)

// Backoff bounds. The first retry comes quickly, because most failures are
// a restart on the other side; the cap is long, because a server that has
// been gone for a day is not coming back in the next minute, and every
// retry costs a connection attempt.
const (
	defaultMinBackoff = 10 * time.Second
	defaultMaxBackoff = 6 * time.Hour
)

// sendTimeout bounds one transaction.
const sendTimeout = 2 * time.Minute

// Sender is the part of the federation client the queue uses.
type Sender interface {
	SendTransaction(ctx context.Context, destination, txnID string, txn federation.Transaction) (federation.SendTransactionResponse, error)
}

// RoomServers says which servers are in a room.
type RoomServers interface {
	JoinedServers(ctx context.Context, roomID string) ([]string, error)
}

// QueueOptions configure a Queue.
type QueueOptions struct {
	ServerName string
	Store      storage.FederationStore
	Sender     Sender
	Rooms      RoomServers
	Policy     *federation.Policy

	MinBackoff time.Duration
	MaxBackoff time.Duration

	Logger *slog.Logger
	Now    func() time.Time
}

// Queue is the outbound queue.
type Queue struct {
	serverName string
	store      storage.FederationStore
	sender     Sender
	rooms      RoomServers
	policy     *federation.Policy
	minBackoff time.Duration
	maxBackoff time.Duration
	log        *slog.Logger
	now        func() time.Time

	// txnPrefix makes transaction IDs unique across restarts: a remote
	// server deduplicates by (origin, txn ID), and an ID reused with
	// different contents would be answered from its cache and the events
	// in it never seen.
	txnPrefix string

	mu       sync.Mutex
	running  bool
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
	incoming []roomserver.Notification
	wake     chan struct{}
	workers  map[string]*worker
	txnSeq   uint64
}

// worker sends one destination's queue.
type worker struct {
	destination string
	wake        chan struct{}
}

// NewQueue builds a stopped queue.
func NewQueue(opts QueueOptions) *Queue {
	if opts.Logger == nil {
		opts.Logger = slog.Default()
	}
	if opts.Now == nil {
		opts.Now = time.Now
	}
	if opts.MinBackoff <= 0 {
		opts.MinBackoff = defaultMinBackoff
	}
	if opts.MaxBackoff < opts.MinBackoff {
		opts.MaxBackoff = max(defaultMaxBackoff, opts.MinBackoff)
	}
	if opts.Policy == nil {
		opts.Policy = federation.NewPolicy(nil, nil)
	}
	return &Queue{
		serverName: opts.ServerName,
		store:      opts.Store,
		sender:     opts.Sender,
		rooms:      opts.Rooms,
		policy:     opts.Policy,
		minBackoff: opts.MinBackoff,
		maxBackoff: opts.MaxBackoff,
		log:        opts.Logger,
		now:        opts.Now,
		wake:       make(chan struct{}, 1),
		workers:    map[string]*worker{},
	}
}

// Start begins sending: the dispatcher, and a worker for every destination
// with something left in the queue from before.
func (q *Queue) Start(ctx context.Context) error {
	q.mu.Lock()
	if q.running {
		q.mu.Unlock()
		return nil
	}
	q.ctx, q.cancel = context.WithCancel(context.WithoutCancel(ctx))
	q.running = true
	q.txnPrefix = strconv.FormatInt(q.now().UnixNano(), 36)
	q.mu.Unlock()

	q.wg.Add(1)
	go q.dispatch()

	destinations, err := q.store.QueuedDestinations(ctx)
	if err != nil {
		return fmt.Errorf("read the outbound queue: %w", err)
	}
	for _, d := range destinations {
		q.kick(d)
	}
	q.log.Info("federation queue started", "destinations_with_backlog", len(destinations))
	return nil
}

// Stop halts every send and waits for the workers to finish. What is
// queued stays queued.
func (q *Queue) Stop() {
	q.mu.Lock()
	if !q.running {
		q.mu.Unlock()
		return
	}
	q.running = false
	q.cancel()
	q.incoming = nil
	q.mu.Unlock()

	q.wg.Wait()

	q.mu.Lock()
	q.workers = map[string]*worker{}
	q.mu.Unlock()
	q.log.Info("federation queue stopped")
}

// Running reports whether the queue is sending.
func (q *Queue) Running() bool {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.running
}

// EventPersisted implements roomserver.Notifier. It must not block, and
// does not: it appends and signals.
//
// While the queue is stopped events are dropped rather than held. That is
// what the federation switch promises — no remote traffic while it is off
// — and holding them would turn switching federation back on after a month
// into a month of events delivered in one burst. A remote server that
// cares about the room fetches what it missed when it next hears from it.
func (q *Queue) EventPersisted(_ context.Context, n roomserver.Notification) {
	if n.Origin != "" && !n.Relay {
		return
	}
	q.mu.Lock()
	if !q.running {
		q.mu.Unlock()
		return
	}
	q.incoming = append(q.incoming, n)
	q.mu.Unlock()
	select {
	case q.wake <- struct{}{}:
	default:
	}
}

// dispatch turns persisted events into queue rows.
func (q *Queue) dispatch() {
	defer q.wg.Done()
	for {
		select {
		case <-q.ctx.Done():
			return
		case <-q.wake:
		}
		q.mu.Lock()
		batch := q.incoming
		q.incoming = nil
		q.mu.Unlock()

		for _, n := range batch {
			if err := q.queueEvent(q.ctx, n); err != nil && q.ctx.Err() == nil {
				q.log.Error("could not queue an event for federation",
					"room_id", n.RoomID, "event_id", n.Event.EventID(), "error", err)
			}
		}
	}
}

// queueEvent writes one event's queue rows.
func (q *Queue) queueEvent(ctx context.Context, n roomserver.Notification) error {
	servers, err := q.rooms.JoinedServers(ctx, n.RoomID)
	if err != nil {
		return err
	}
	// A membership event also goes to the server of the user it is about:
	// a server whose last user was just kicked is no longer joined, and is
	// exactly the one that needs to hear it.
	if key, ok := n.Event.StateKey(); ok && n.Event.Type() == event.TypeMember {
		if server, err := id.ServerNameOf(key); err == nil && !slices.Contains(servers, server) {
			servers = append(servers, server)
		}
	}

	var items []storage.QueuedItem
	for _, server := range servers {
		if server == q.serverName || server == n.Origin || !q.policy.Allowed(server) {
			continue
		}
		items = append(items, storage.QueuedItem{
			Destination: server,
			Kind:        storage.QueueKindPDU,
			RoomID:      n.RoomID,
			EventID:     n.Event.EventID(),
			Payload:     n.Event.JSON(),
			QueuedAtMS:  q.now().UnixMilli(),
		})
	}
	return q.enqueue(ctx, items)
}

// QueueEDU queues an EDU for a set of destinations.
func (q *Queue) QueueEDU(ctx context.Context, destinations []string, edu federation.EDU) {
	if !q.Running() {
		return
	}
	payload, err := json.Marshal(edu)
	if err != nil {
		q.log.Error("could not encode an EDU", "type", edu.Type, "error", err)
		return
	}
	var items []storage.QueuedItem
	for _, server := range destinations {
		if server == q.serverName || !q.policy.Allowed(server) {
			continue
		}
		items = append(items, storage.QueuedItem{
			Destination: server,
			Kind:        storage.QueueKindEDU,
			Payload:     payload,
			QueuedAtMS:  q.now().UnixMilli(),
		})
	}
	if err := q.enqueue(ctx, items); err != nil {
		q.log.Error("could not queue an EDU", "type", edu.Type, "error", err)
	}
}

// enqueue writes rows and wakes their workers.
func (q *Queue) enqueue(ctx context.Context, items []storage.QueuedItem) error {
	if len(items) == 0 {
		return nil
	}
	if _, err := q.store.Enqueue(ctx, items); err != nil {
		return err
	}
	seen := map[string]bool{}
	for _, item := range items {
		if !seen[item.Destination] {
			seen[item.Destination] = true
			q.kick(item.Destination)
		}
	}
	return nil
}

// kick wakes a destination's worker, starting one if there is none.
func (q *Queue) kick(destination string) {
	q.mu.Lock()
	defer q.mu.Unlock()
	if !q.running {
		return
	}
	w, ok := q.workers[destination]
	if !ok {
		w = &worker{destination: destination, wake: make(chan struct{}, 1)}
		q.workers[destination] = w
		q.wg.Add(1)
		go q.run(w)
	}
	select {
	case w.wake <- struct{}{}:
	default:
	}
}

// run is one destination's worker. It sends until the queue for the
// destination is empty, then exits; the next enqueue starts a new one.
func (q *Queue) run(w *worker) {
	defer q.wg.Done()
	for {
		dest, err := q.store.Destination(q.ctx, w.destination)
		if err != nil {
			if q.ctx.Err() != nil {
				return
			}
			q.log.Error("could not read a destination's backoff", "destination", w.destination, "error", err)
			dest = storage.Destination{ServerName: w.destination}
		}
		if wait := time.Duration(dest.RetryAtMS-q.now().UnixMilli()) * time.Millisecond; wait > 0 {
			// A wake-up during the wait re-reads the backoff, which is how
			// an administrator's reset takes effect at once; an ordinary
			// enqueue finds the backoff unchanged and waits again.
			timer := time.NewTimer(wait)
			select {
			case <-q.ctx.Done():
				timer.Stop()
				return
			case <-w.wake:
				timer.Stop()
				continue
			case <-timer.C:
			}
		}

		sent, err := q.sendOnce(w.destination, dest)
		if q.ctx.Err() != nil {
			return
		}
		if err != nil {
			// A failed send has recorded its backoff, which the next pass
			// waits out. The pause here is for the other kind of failure —
			// the queue's own database — which records nothing and would
			// otherwise be retried in a tight loop.
			select {
			case <-q.ctx.Done():
				return
			case <-time.After(time.Second):
			}
			continue
		}
		if !sent && q.retire(w) {
			return
		}
	}
}

// retire removes an idle worker, unless something was queued for it in
// the meantime.
func (q *Queue) retire(w *worker) bool {
	q.mu.Lock()
	defer q.mu.Unlock()
	select {
	case <-w.wake:
		return false
	default:
	}
	if q.workers[w.destination] == w {
		delete(q.workers, w.destination)
	}
	return true
}

// sendOnce sends one transaction. It reports whether there was anything to
// send.
func (q *Queue) sendOnce(destination string, health storage.Destination) (bool, error) {
	pdus, err := q.store.Queued(q.ctx, destination, storage.QueueKindPDU, maxPDUsPerTxn)
	if err != nil {
		return false, err
	}
	edus, err := q.store.Queued(q.ctx, destination, storage.QueueKindEDU, maxEDUsPerTxn)
	if err != nil {
		return false, err
	}
	if len(pdus) == 0 && len(edus) == 0 {
		return false, nil
	}

	// A destination the policy no longer allows gets nothing, and what
	// was queued for it is dropped: the allow and deny lists can change at
	// runtime, and a backlog for a server that was just denied is exactly
	// what must not go out.
	if !q.policy.Allowed(destination) {
		return true, q.store.Dequeue(q.ctx, append(itemIDs(pdus), itemIDs(edus)...))
	}

	txn := federation.Transaction{
		Origin:         q.serverName,
		OriginServerTS: q.now().UnixMilli(),
		PDUs:           make([]json.RawMessage, 0, len(pdus)),
	}
	for _, item := range pdus {
		txn.PDUs = append(txn.PDUs, item.Payload)
	}
	for _, item := range edus {
		var edu federation.EDU
		if err := json.Unmarshal(item.Payload, &edu); err == nil {
			txn.EDUs = append(txn.EDUs, edu)
		}
	}

	ctx, cancel := context.WithTimeout(q.ctx, sendTimeout)
	defer cancel()
	resp, err := q.sender.SendTransaction(ctx, destination, q.nextTxnID(), txn)
	if err != nil {
		if q.ctx.Err() != nil {
			return true, err
		}
		q.recordFailure(destination, health, err)
		return true, err
	}

	for eventID, result := range resp.PDUs {
		if result.Error != "" {
			q.log.Info("a destination refused an event",
				"destination", destination, "event_id", eventID, "error", result.Error)
		}
	}
	if err := q.store.Dequeue(q.ctx, append(itemIDs(pdus), itemIDs(edus)...)); err != nil {
		return true, err
	}
	if health.Failures > 0 || health.LastSuccessMS == 0 {
		q.log.Info("destination is reachable", "destination", destination, "after_failures", health.Failures)
		if err := q.store.SetDestination(q.ctx, storage.Destination{
			ServerName:    destination,
			LastSuccessMS: q.now().UnixMilli(),
			LastFailureMS: health.LastFailureMS,
			LastError:     health.LastError,
		}); err != nil {
			q.log.Warn("could not record a destination's recovery", "destination", destination, "error", err)
		}
	}
	return true, nil
}

// recordFailure backs a destination off.
func (q *Queue) recordFailure(destination string, health storage.Destination, err error) {
	failures := health.Failures + 1
	backoff := q.backoff(failures)
	now := q.now()
	q.log.Info("sending to a destination failed",
		"destination", destination, "failures", failures, "retry_in", backoff, "error", err)
	if serr := q.store.SetDestination(q.ctx, storage.Destination{
		ServerName:    destination,
		Failures:      failures,
		RetryAtMS:     now.Add(backoff).UnixMilli(),
		LastSuccessMS: health.LastSuccessMS,
		LastFailureMS: now.UnixMilli(),
		LastError:     truncate(err.Error(), 500),
	}); serr != nil {
		q.log.Warn("could not record a destination's backoff", "destination", destination, "error", serr)
	}
}

// backoff is the wait after a number of consecutive failures: doubling
// from the minimum up to the cap, with a fifth of jitter so that a server
// coming back is not hit by every other server's retry at once.
func (q *Queue) backoff(failures int) time.Duration {
	d := q.minBackoff
	for i := 1; i < failures && d < q.maxBackoff; i++ {
		d *= 2
	}
	if d > q.maxBackoff {
		d = q.maxBackoff
	}
	jitter := time.Duration(rand.Int64N(int64(d)/5 + 1))
	return d + jitter
}

// nextTxnID mints a transaction ID.
func (q *Queue) nextTxnID() string {
	q.mu.Lock()
	defer q.mu.Unlock()
	q.txnSeq++
	return q.txnPrefix + "_" + strconv.FormatUint(q.txnSeq, 10)
}

// Destinations lists every destination the queue knows, for the admin API.
func (q *Queue) Destinations(ctx context.Context) ([]storage.DestinationSummary, error) {
	return q.store.Destinations(ctx)
}

// ResetBackoff clears a destination's backoff so that it is tried now,
// for an administrator who knows it is back.
func (q *Queue) ResetBackoff(ctx context.Context, destination string) error {
	health, err := q.store.Destination(ctx, destination)
	if err != nil {
		return err
	}
	health.Failures = 0
	health.RetryAtMS = 0
	if err := q.store.SetDestination(ctx, health); err != nil {
		return err
	}
	q.kick(destination)
	return nil
}

func itemIDs(items []storage.QueuedItem) []int64 {
	out := make([]int64, 0, len(items))
	for _, item := range items {
		out = append(out, item.ID)
	}
	return out
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n]
}