package jetstream

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"tangled.org/core/log"
)

const (
	gapRewind = time.Second
	// shutdown saves the cursor after this even if events are still held
	heldGrace = 30 * time.Second
)

type DB interface {
	GetLastTimeUs() (int64, error)
	SaveLastTimeUs(int64) error
}

type Set[T comparable] map[T]struct{}

type JetstreamClient struct {
	cfg StreamConfig
	l   *slog.Logger

	logDids         bool
	wantedDids      Set[string]
	unfilteredNsids Set[string]
	db              DB
	waitForDid      bool
	mu              sync.RWMutex

	lastSeenUs atomic.Int64
	marks      watermark
	holding    atomic.Int64

	// the first event of each connection clears it, so the gap check runs once
	anchorUs atomic.Int64
	onGap    func(context.Context)

	done chan struct{}

	saveEvery time.Duration
}

func (j *JetstreamClient) AddDid(did string) {
	if did == "" {
		return
	}

	if j.logDids {
		j.l.Info("adding did to in-memory filter", "did", did)
	}
	j.mu.Lock()
	j.wantedDids[did] = struct{}{}
	j.mu.Unlock()
}

// call before StartJetstream, fn runs on the reader so don't block in it
func (j *JetstreamClient) SetGapHandler(fn func(context.Context)) {
	j.onGap = fn
}

func (j *JetstreamClient) ExemptCollection(nsid string) {
	j.mu.Lock()
	j.unfilteredNsids[nsid] = struct{}{}
	j.mu.Unlock()
}

func (j *JetstreamClient) RemoveDid(did string) {
	if did == "" {
		return
	}

	if j.logDids {
		j.l.Info("removing did from in-memory filter", "did", did)
	}
	j.mu.Lock()
	delete(j.wantedDids, did)
	j.mu.Unlock()
}

type processor func(context.Context, *Event) error

func (j *JetstreamClient) withDidFilter(processFunc processor) processor {
	// since this closure references j.WantedDids; it should auto-update
	// existing instances of the closure when j.WantedDids is mutated
	return func(ctx context.Context, evt *Event) error {
		if anchor := j.anchorUs.Swap(0); anchor != 0 && evt.TimeUS > anchor && j.onGap != nil {
			j.l.Warn("jetstream history no longer reaches the last event seen", "last", anchor, "first", evt.TimeUS)
			j.onGap(ctx)
		}

		ticket := j.ticket(evt.TimeUS)
		h := &hold{ticket: ticket}
		ctx = context.WithValue(ctx, holdKey{}, h)

		j.mu.RLock()
		// empty filter => all dids allowed
		matches := len(j.wantedDids) == 0
		if !matches {
			if _, ok := j.wantedDids[evt.Did]; ok {
				matches = true
			}
		}
		if !matches && evt.Commit != nil {
			if _, ok := j.unfilteredNsids[evt.Commit.Collection]; ok {
				matches = true
			}
		}
		j.mu.RUnlock()

		var err error
		if matches {
			err = processFunc(ctx, evt)
		}
		if err != nil && ctx.Err() != nil {
			// shutdown cut this event short, so the cursor stays before it and the next start replays it
			return err
		}

		if !h.held.Load() {
			ticket.Done()
		}
		return err
	}
}

// StartJetstream replaces cfg.Cursor with the db one, so don't set it
func NewJetstreamClient(endpoint string, collections []string, cfg *StreamConfig, logger *slog.Logger, db DB, waitForDid, logDids bool) (*JetstreamClient, error) {
	if cfg == nil {
		cfg = &StreamConfig{Endpoint: endpoint, Collections: collections, Compress: true}
	}

	return &JetstreamClient{
		cfg:             *cfg,
		db:              db,
		l:               logger,
		wantedDids:      make(map[string]struct{}),
		unfilteredNsids: make(map[string]struct{}),

		logDids: logDids,

		// This will make the goroutine in StartJetstream wait until
		// j.wantedDids has been populated, typically using addDids.
		waitForDid: waitForDid,

		done: make(chan struct{}),

		saveEvery: time.Minute,
	}, nil
}

// StartJetstream starts the jetstream client and processes events using the provided processFunc.
// The client persists the last time_us cursor itself via the DB it was constructed with.
func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *Event) error) error {
	cfg := j.cfg
	cfg.TailWhenNoCursor = true
	cfg.Cursor = func() (int64, bool, error) { return j.cursor(ctx) }
	stream, err := NewStream(cfg, j.l)
	if err != nil {
		return fmt.Errorf("failed to create jetstream client: %w", err)
	}

	handle := j.withDidFilter(processFunc)
	go func() {
		defer close(j.done)
		if j.waitForDid && !j.waitForWantedDid(ctx) {
			return
		}
		j.l.Info("done waiting for did")

		ticked := make(chan struct{})
		go func() {
			defer close(ticked)
			j.periodicLastTimeSave(ctx)
		}()

		// returning the error would replay it forever, so log and skip
		stream.Run(ctx, func(ctx context.Context, e *Event) error {
			if err := handle(ctx, e); err != nil {
				if ctx.Err() != nil {
					return err
				}
				j.l.Error("error processing event", "error", err, "time_us", e.TimeUS)
			}
			return nil
		})
		<-ticked
		j.waitForHeld()
		// save last, nothing may move the cursor once it runs
		j.saveLastSeen(ctx)
	}()

	return nil
}

func (j *JetstreamClient) waitForHeld() {
	for deadline := time.Now().Add(heldGrace); j.holding.Load() > 0 && time.Now().Before(deadline); {
		time.Sleep(10 * time.Millisecond)
	}
}

// one past the last finished event, not the event itself
func (j *JetstreamClient) LastSeenUs() int64 {
	return j.lastSeenUs.Load()
}

// wait on it before closing the db, the final cursor save needs it open
// it never closes if StartJetstream wasn't called
func (j *JetstreamClient) Done() <-chan struct{} {
	return j.done
}

func (j *JetstreamClient) waitForWantedDid(ctx context.Context) bool {
	for {
		j.mu.RLock()
		hasDid := len(j.wantedDids) != 0
		j.mu.RUnlock()
		if hasDid {
			return true
		}
		select {
		case <-ctx.Done():
			return false
		case <-time.After(time.Second):
		}
	}
}

// save cursor periodically
func (j *JetstreamClient) periodicLastTimeSave(ctx context.Context) {
	ticker := time.NewTicker(j.saveEvery)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			j.saveLastSeen(ctx)
		}
	}
}

func (j *JetstreamClient) resume(ctx context.Context) *int64 {
	cursor := j.resumeCursor(ctx)
	j.anchorUs.Store(0)
	if cursor == nil || j.onGap == nil || *cursor <= 0 {
		return cursor
	}
	// cursor sits one past its event, hence the -1
	j.anchorUs.Store(*cursor - 1)
	rewound := max(*cursor-gapRewind.Microseconds(), 0)
	return &rewound
}

func (j *JetstreamClient) cursor(ctx context.Context) (int64, bool, error) {
	cursor := j.resume(ctx)
	if cursor == nil {
		return 0, false, nil
	}
	return *cursor, true, nil
}

func (j *JetstreamClient) resumeCursor(ctx context.Context) *int64 {
	if seen := j.lastSeenUs.Load(); seen != 0 {
		return &seen
	}
	return j.getLastTimeUs(ctx)
}

// the saved cursor is always an event time_us, never a wall clock
func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 {
	l := log.FromContext(ctx)
	lastTimeUs, err := j.db.GetLastTimeUs()
	if err != nil {
		l.Warn("couldn't get last time us, starting from now", "error", err)
		return nil
	}

	l.Info("found last time_us", "time_us", lastTimeUs)
	return &lastTimeUs
}

func (j *JetstreamClient) saveLastSeen(ctx context.Context) {
	if seen := j.lastSeenUs.Load(); seen != 0 {
		if err := j.db.SaveLastTimeUs(seen); err != nil {
			log.FromContext(ctx).Error("failed to save cursor", "error", err)
		}
	}
}