// Package jetstream consumes real-time AT Protocol events from a Jetstream relay.
//
// Jetstream is a WebSocket-based relay that delivers a filtered stream of AT
// Protocol repository events (commits, identity changes, account changes).
// This package handles connection management, reconnection with backoff,
// endpoint rotation, cursor tracking, optional zstd decompression, and
// runtime-mutable subscription filters via the jetstream options_update
// control frame.
//
// Basic usage:
//
//	consumer := jetstream.New(&jetstream.Config{
//	    WantedCollections: []string{"app.bsky.feed.post"},
//	}, func(ctx context.Context, evt *jetstream.Event) error {
//	    fmt.Printf("new post from %s\n", evt.DID)
//	    return nil
//	})
//	consumer.Start(ctx)
//	defer consumer.Stop()
package jetstream

import (
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"net/url"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
)

// zstdDictionary is the shared dictionary the Jetstream relay compresses
// each frame against. Sourced from:
//
//	https://github.com/bluesky-social/jetstream/raw/main/pkg/models/zstd_dictionary
//
//go:embed zstd_dictionary
var zstdDictionary []byte

// DefaultEndpoints are the public Jetstream relay endpoints.
var DefaultEndpoints = []string{
	"wss://jetstream1.us-east.bsky.network/subscribe",
	"wss://jetstream2.us-east.bsky.network/subscribe",
	"wss://jetstream1.us-west.bsky.network/subscribe",
	"wss://jetstream2.us-west.bsky.network/subscribe",
}

// Default liveness settings. Jetstream's server-side filters mean a connection
// can sit silent indefinitely if no matching events are flowing. Application
// pings keep the socket warm and let callers distinguish "upstream has nothing
// for us" from "upstream is dead". The read deadline must be comfortably
// greater than 2× the ping interval so a single dropped pong doesn't trip it.
const (
	defaultPingInterval     = 30 * time.Second
	defaultReadDeadline     = 90 * time.Second
	defaultPingWriteTimeout = 10 * time.Second
)

// Event is a single event from the Jetstream relay.
type Event struct {
	DID      string    `json:"did"`
	TimeUS   int64     `json:"time_us"`
	Kind     string    `json:"kind"` // "commit", "identity", "account"
	Commit   *Commit   `json:"commit,omitempty"`
	Identity *Identity `json:"identity,omitempty"`
	Account  *Account  `json:"account,omitempty"`
}

// Commit is the commit payload within an Event.
type Commit struct {
	Rev        string          `json:"rev"`
	Operation  string          `json:"operation"` // "create", "update", "delete"
	Collection string          `json:"collection"`
	RKey       string          `json:"rkey"`
	Record     json.RawMessage `json:"record,omitempty"`
	CID        string          `json:"cid"`
}

// Identity is the identity payload of an Event. Emitted when a DID's handle
// or PDS endpoint changes.
type Identity struct {
	DID    string `json:"did"`
	Handle string `json:"handle,omitempty"`
	Seq    int64  `json:"seq"`
	Time   string `json:"time"`
}

// Account is the account payload of an Event. Status is one of "active",
// "deleted", "deactivated", "suspended", "takendown".
type Account struct {
	Active bool   `json:"active"`
	DID    string `json:"did"`
	Seq    int64  `json:"seq"`
	Status string `json:"status,omitempty"`
	Time   string `json:"time"`
}

// Handler is called for each event received from Jetstream.
// Returning an error logs a warning but does not stop the consumer.
type Handler func(ctx context.Context, event *Event) error

// CursorStore persists the Jetstream cursor across restarts.
// If nil, the cursor is tracked in memory only (replay from live on restart).
type CursorStore interface {
	GetCursor(ctx context.Context) (int64, error)
	SetCursor(ctx context.Context, cursor int64) error
}

// Config configures a Jetstream consumer.
type Config struct {
	// Endpoints is the list of Jetstream WebSocket URLs. Defaults to DefaultEndpoints.
	Endpoints []string

	// WantedCollections filters events to specific NSIDs.
	// Empty means all collections (high volume).
	WantedCollections []string

	// WantedDIDs filters events to specific DID subjects.
	// Empty means all DIDs. When combined with WantedCollections, the
	// jetstream relay applies AND semantics — events must match both filters.
	WantedDIDs []string

	// Compress enables zstd compression on the wire. The Jetstream zstd
	// dictionary is embedded in the package and applied automatically.
	Compress bool

	// CursorStore persists the cursor for resume after restart.
	// If nil, the consumer starts from live on each restart.
	CursorStore CursorStore

	// CursorPersistEvery controls how often the cursor is flushed to CursorStore.
	// Defaults to every 1000 events.
	CursorPersistEvery int64

	// PingInterval is how often application-level ping frames are sent to keep
	// the socket warm. Defaults to 30s. Set to a negative value to disable
	// pings entirely (rely on the read deadline alone).
	PingInterval time.Duration

	// ReadDeadline is the per-read websocket deadline, refreshed on every
	// successful read and every pong. Should be comfortably greater than 2×
	// PingInterval. Defaults to 90s.
	ReadDeadline time.Duration

	// OnConnect is called each time a WebSocket connection is established.
	// Call Consumer.CurrentEndpoint to learn which endpoint we connected to.
	OnConnect func()

	// OnDisconnect is called each time a connection is lost.
	OnDisconnect func()

	// OnError is called when the handler returns an error.
	// If nil, errors are silently dropped (caller should log in the handler).
	OnError func(err error, event *Event)
}

func (c *Config) endpoints() []string {
	if len(c.Endpoints) > 0 {
		return c.Endpoints
	}
	return DefaultEndpoints
}

func (c *Config) cursorPersistEvery() int64 {
	if c.CursorPersistEvery > 0 {
		return c.CursorPersistEvery
	}
	return 1000
}

// pingInterval returns the configured ping interval, the default if zero, or
// zero if pings are explicitly disabled (negative config value).
func (c *Config) pingInterval() time.Duration {
	if c.PingInterval < 0 {
		return 0
	}
	if c.PingInterval == 0 {
		return defaultPingInterval
	}
	return c.PingInterval
}

func (c *Config) readDeadline() time.Duration {
	if c.ReadDeadline > 0 {
		return c.ReadDeadline
	}
	return defaultReadDeadline
}

// optionsUpdateMessage is the jetstream control frame for runtime filter
// updates. See https://github.com/bluesky-social/jetstream for the schema.
type optionsUpdateMessage struct {
	Type    string               `json:"type"`
	Payload optionsUpdatePayload `json:"payload"`
}

type optionsUpdatePayload struct {
	WantedCollections []string `json:"wantedCollections"`
	WantedDIDs        []string `json:"wantedDids"`
}

// Consumer consumes events from a Jetstream relay.
type Consumer struct {
	cfg     *Config
	handler Handler

	conn               *websocket.Conn
	connMu             sync.Mutex
	currentEndpointIdx int
	currentEndpoint    atomic.Value // string

	// writeMu serializes control-frame writes (pings and options_update).
	// gorilla/websocket is not safe for concurrent writes.
	writeMu sync.Mutex

	// Mutable subscription state. Mirrors the config at construction time but
	// can be replaced at runtime via UpdateOptions.
	optMu             sync.RWMutex
	wantedCollections []string
	wantedDIDs        []string

	zstdDecoder *zstd.Decoder

	cursor         atomic.Int64
	eventsReceived atomic.Int64
	bytesReceived  atomic.Int64
	connected      atomic.Bool

	stopCh chan struct{}
	wg     sync.WaitGroup
}

// New creates a new Consumer. Call Start to begin consuming events.
func New(cfg *Config, handler Handler) *Consumer {
	opts := []zstd.DOption{zstd.WithDecoderConcurrency(1)}
	if cfg.Compress {
		opts = append(opts, zstd.WithDecoderDicts(zstdDictionary))
	}
	decoder, err := zstd.NewReader(nil, opts...)
	if err != nil {
		// zstd.NewReader with nil src only fails on bad options
		panic(fmt.Sprintf("jetstream: create zstd decoder: %v", err))
	}

	c := &Consumer{
		cfg:               cfg,
		handler:           handler,
		stopCh:            make(chan struct{}),
		zstdDecoder:       decoder,
		wantedCollections: slices.Clone(cfg.WantedCollections),
		wantedDIDs:        slices.Clone(cfg.WantedDIDs),
	}
	c.currentEndpoint.Store("")

	if cfg.CursorStore != nil {
		if cursor, err := cfg.CursorStore.GetCursor(context.Background()); err == nil && cursor > 0 {
			c.cursor.Store(cursor)
		}
	}

	return c
}

// Start begins consuming events in a background goroutine.
func (c *Consumer) Start(ctx context.Context) {
	c.wg.Go(func() {
		c.run(ctx)
	})
}

// Stop gracefully shuts down the consumer and waits for it to finish.
func (c *Consumer) Stop() {
	close(c.stopCh)
	c.connMu.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connMu.Unlock()
	c.wg.Wait()
	c.zstdDecoder.Close()
}

// IsConnected reports whether the consumer is currently connected.
func (c *Consumer) IsConnected() bool {
	return c.connected.Load()
}

// CurrentEndpoint returns the endpoint URL of the active connection, or the
// empty string if not connected. Safe for concurrent use.
func (c *Consumer) CurrentEndpoint() string {
	v, _ := c.currentEndpoint.Load().(string)
	return v
}

// Stats returns cumulative event and byte counts since Start was called.
func (c *Consumer) Stats() (eventsReceived, bytesReceived int64) {
	return c.eventsReceived.Load(), c.bytesReceived.Load()
}

// UpdateOptions atomically replaces the subscription filters. If currently
// connected, an options_update control frame is sent to the relay so the new
// filters take effect immediately. Otherwise the new filters apply on the
// next reconnect via the WebSocket URL.
//
// Passing nil slices is equivalent to passing empty slices: the filter is
// cleared and the relay delivers all events for that dimension.
func (c *Consumer) UpdateOptions(wantedCollections, wantedDIDs []string) error {
	collsCopy := slices.Clone(wantedCollections)
	didsCopy := slices.Clone(wantedDIDs)

	c.optMu.Lock()
	c.wantedCollections = collsCopy
	c.wantedDIDs = didsCopy
	c.optMu.Unlock()

	c.connMu.Lock()
	conn := c.conn
	c.connMu.Unlock()
	if conn == nil {
		return nil // applied via URL on next reconnect
	}

	msg := optionsUpdateMessage{
		Type: "options_update",
		Payload: optionsUpdatePayload{
			WantedCollections: collsCopy,
			WantedDIDs:        didsCopy,
		},
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal options_update: %w", err)
	}

	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
		return fmt.Errorf("send options_update: %w", err)
	}
	return nil
}

func (c *Consumer) run(ctx context.Context) {
	backoff := time.Second
	const maxBackoff = 30 * time.Second

	for {
		select {
		case <-ctx.Done():
			return
		case <-c.stopCh:
			return
		default:
		}

		endpoints := c.cfg.endpoints()
		endpoint := endpoints[c.currentEndpointIdx]

		if err := c.connectAndConsume(ctx, endpoint); err != nil {
			c.connected.Store(false)
			if c.cfg.OnDisconnect != nil {
				c.cfg.OnDisconnect()
			}

			// Rotate to next endpoint
			c.currentEndpointIdx = (c.currentEndpointIdx + 1) % len(endpoints)

			select {
			case <-ctx.Done():
				return
			case <-c.stopCh:
				return
			case <-time.After(backoff):
			}

			backoff *= 2
			if backoff > maxBackoff {
				backoff = maxBackoff
			}
		} else {
			backoff = time.Second
		}
	}
}

func (c *Consumer) connectAndConsume(ctx context.Context, endpoint string) error {
	wsURL, err := c.buildURL(endpoint)
	if err != nil {
		return fmt.Errorf("build URL: %w", err)
	}

	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
	conn, _, err := dialer.DialContext(ctx, wsURL, nil)
	if err != nil {
		return fmt.Errorf("dial: %w", err)
	}

	c.connMu.Lock()
	c.conn = conn
	c.connMu.Unlock()

	c.currentEndpoint.Store(endpoint)
	c.connected.Store(true)
	if c.cfg.OnConnect != nil {
		c.cfg.OnConnect()
	}

	defer func() {
		c.connMu.Lock()
		if c.conn != nil {
			c.conn.Close()
			c.conn = nil
		}
		c.connMu.Unlock()
		c.currentEndpoint.Store("")
		c.connected.Store(false)
	}()

	deadline := c.cfg.readDeadline()
	conn.SetReadDeadline(time.Now().Add(deadline))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(deadline))
	})

	// Pings run in their own goroutine. Cancel them before the connection
	// teardown above so we don't write to a closed conn.
	pingCtx, cancelPing := context.WithCancel(ctx)
	defer cancelPing()
	c.startPingLoop(pingCtx, conn)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-c.stopCh:
			return nil
		default:
		}

		_, msg, err := conn.ReadMessage()
		if err != nil {
			return fmt.Errorf("read: %w", err)
		}

		// Refresh deadline on every successful read. The PongHandler covers
		// idle periods; this covers active ones.
		conn.SetReadDeadline(time.Now().Add(deadline))

		c.bytesReceived.Add(int64(len(msg)))

		if err := c.process(ctx, msg); err != nil {
			if c.cfg.OnError != nil {
				// We don't have the event here since parsing may have failed,
				// pass nil to signal a parse/process error
				c.cfg.OnError(err, nil)
			}
		}
	}
}

func (c *Consumer) startPingLoop(ctx context.Context, conn *websocket.Conn) {
	interval := c.cfg.pingInterval()
	if interval <= 0 {
		return
	}
	go func() {
		t := time.NewTicker(interval)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-c.stopCh:
				return
			case <-t.C:
				c.writeMu.Lock()
				err := conn.WriteControl(
					websocket.PingMessage,
					nil,
					time.Now().Add(defaultPingWriteTimeout),
				)
				c.writeMu.Unlock()
				if err != nil {
					// The read loop will surface the underlying connection
					// failure; we just stop pinging.
					return
				}
			}
		}
	}()
}

func (c *Consumer) buildURL(endpoint string) (string, error) {
	u, err := url.Parse(endpoint)
	if err != nil {
		return "", err
	}

	q := u.Query()

	c.optMu.RLock()
	for _, coll := range c.wantedCollections {
		q.Add("wantedCollections", coll)
	}
	for _, did := range c.wantedDIDs {
		q.Add("wantedDids", did)
	}
	c.optMu.RUnlock()

	if c.cfg.Compress {
		q.Set("compress", "true")
	}
	if cursor := c.cursor.Load(); cursor > 0 {
		// Rewind 5 seconds to cover any gaps at reconnect
		rewind := cursor - (5 * time.Second.Microseconds())
		q.Set("cursor", fmt.Sprintf("%d", rewind))
	}

	u.RawQuery = q.Encode()
	return u.String(), nil
}

func (c *Consumer) process(ctx context.Context, data []byte) error {
	if c.cfg.Compress {
		decompressed, err := c.zstdDecoder.DecodeAll(data, nil)
		if err != nil {
			return fmt.Errorf("decompress: %w", err)
		}
		data = decompressed
	}

	var event Event
	if err := json.Unmarshal(data, &event); err != nil {
		return fmt.Errorf("unmarshal event: %w", err)
	}

	c.eventsReceived.Add(1)

	if event.TimeUS > 0 {
		c.cursor.Store(event.TimeUS)

		if c.cfg.CursorStore != nil && c.eventsReceived.Load()%c.cfg.cursorPersistEvery() == 0 {
			if err := c.cfg.CursorStore.SetCursor(ctx, event.TimeUS); err != nil {
				// Non-fatal: log via OnError if configured
				if c.cfg.OnError != nil {
					c.cfg.OnError(fmt.Errorf("persist cursor: %w", err), nil)
				}
			}
		}
	}

	if err := c.handler(ctx, &event); err != nil {
		if c.cfg.OnError != nil {
			c.cfg.OnError(err, &event)
		}
	}

	return nil
}