at-container-registry/pkg/appview/jetstream/worker.go

// Package jetstream provides an ATProto Jetstream consumer for real-time updates.
// It connects to the Bluesky Jetstream WebSocket, processes repository events,
// indexes manifests and tags, and populates the AppView database for the web UI.
package jetstream

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/url"
	"sync"
	"time"

	"atcr.io/pkg/appview/db"
	"atcr.io/pkg/atproto"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
)

// UserCache caches DID -> handle/PDS mappings to avoid repeated lookups
type UserCache struct {
	cache map[string]*db.User
}

// EventCallback is called for each processed event
type EventCallback func(timeUS int64)

// Worker consumes Jetstream events and populates the UI database
type Worker struct {
	db                   *sql.DB
	jetstreamURL         string
	startCursor          int64
	wantedCollections    []string
	debugCollectionCount int
	processor            *Processor // Shared processor for DB operations
	eventCallback        EventCallback
	connStartTime        time.Time // Track when connection started for debugging

	// Ping/pong tracking for connection health
	pingsSent     int64
	pongsReceived int64
	lastPongTime  time.Time
	pongMutex     sync.Mutex

	// In-memory cursor tracking for reconnects
	lastCursor  int64
	cursorMutex sync.RWMutex
}

// NewWorker creates a new Jetstream worker
// startCursor: Unix microseconds timestamp to start from (0 = start from now)
func NewWorker(database *sql.DB, jetstreamURL string, startCursor int64) *Worker {
	if jetstreamURL == "" {
		jetstreamURL = "wss://jetstream2.us-west.bsky.network/subscribe"
	}

	return &Worker{
		db:           database,
		jetstreamURL: jetstreamURL,
		startCursor:  startCursor,
		wantedCollections: []string{
			"io.atcr.*", // Subscribe to all ATCR collections
		},
		processor: NewProcessor(database, true), // Use cache for live streaming
	}
}

// Start begins consuming Jetstream events
// This is a blocking function that runs until the context is cancelled
func (w *Worker) Start(ctx context.Context) error {
	// Build connection URL with filters
	u, err := url.Parse(w.jetstreamURL)
	if err != nil {
		return fmt.Errorf("invalid jetstream URL: %w", err)
	}

	q := u.Query()
	for _, collection := range w.wantedCollections {
		q.Add("wantedCollections", collection)
	}

	// Add cursor if specified (for backfilling historical data or reconnects)
	if w.startCursor > 0 {
		q.Set("cursor", fmt.Sprintf("%d", w.startCursor))

		// Calculate lag (cursor is in microseconds)
		now := time.Now().UnixMicro()
		lagSeconds := float64(now-w.startCursor) / 1_000_000.0
		slog.Info("Jetstream starting from cursor", "cursor", w.startCursor, "lag_seconds", lagSeconds)
	}

	// Disable compression for now to debug
	// q.Set("compress", "true")
	u.RawQuery = q.Encode()

	slog.Info("Connecting to Jetstream", "url", u.String())

	// Connect to Jetstream
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, u.String(), nil)
	if err != nil {
		return fmt.Errorf("failed to connect to jetstream: %w", err)
	}
	defer conn.Close()

	// Track connection start time for debugging
	w.connStartTime = time.Now()

	// Reset ping/pong counters for this connection
	w.pongMutex.Lock()
	w.pingsSent = 0
	w.pongsReceived = 0
	w.lastPongTime = time.Now()
	w.pongMutex.Unlock()

	// Set up pong handler - called when server responds to our ping
	conn.SetPongHandler(func(appData string) error {
		w.pongMutex.Lock()
		w.pongsReceived++
		w.lastPongTime = time.Now()
		w.pongMutex.Unlock()

		// Reset read deadline - we know connection is alive
		// Allow 90 seconds for next pong (3x ping interval)
		conn.SetReadDeadline(time.Now().Add(90 * time.Second))
		return nil
	})

	// Set initial read deadline
	conn.SetReadDeadline(time.Now().Add(90 * time.Second))

	// Create zstd decoder for decompressing messages
	decoder, err := zstd.NewReader(nil)
	if err != nil {
		return fmt.Errorf("failed to create zstd decoder: %w", err)
	}
	defer decoder.Close()

	slog.Info("Connected to Jetstream, listening for events...")

	// Start heartbeat ticker to show Jetstream is alive
	heartbeatTicker := time.NewTicker(30 * time.Second)
	defer heartbeatTicker.Stop()

	// Start ping ticker for keepalive
	pingTicker := time.NewTicker(30 * time.Second)
	defer pingTicker.Stop()

	// Start ping sender goroutine
	pingDone := make(chan struct{})
	defer close(pingDone)

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-pingDone:
				return
			case <-pingTicker.C:
				// Check if we've received a pong recently
				w.pongMutex.Lock()
				timeSinceLastPong := time.Since(w.lastPongTime)
				pingsTotal := w.pingsSent
				pongsTotal := w.pongsReceived
				w.pongMutex.Unlock()

				// If no pong for 60 seconds, connection is likely dead
				if timeSinceLastPong > 60*time.Second {
					slog.Info("Jetstream no pong received, closing connection", "time_since_last_pong", timeSinceLastPong, "pings_sent", pingsTotal, "pongs_received", pongsTotal)
					conn.Close()
					return
				}

				// Send ping with write deadline
				conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
					slog.Warn("Jetstream failed to send ping", "error", err)
					conn.Close()
					return
				}

				w.pongMutex.Lock()
				w.pingsSent++
				w.pongMutex.Unlock()
			}
		}
	}()

	eventCount := 0
	lastHeartbeat := time.Now()

	// Read messages
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-heartbeatTicker.C:
			elapsed := time.Since(lastHeartbeat)
			slog.Info("Jetstream alive", "events_processed", eventCount, "elapsed_seconds", elapsed.Seconds())
			eventCount = 0
			lastHeartbeat = time.Now()
		default:
			_, message, err := conn.ReadMessage()
			if err != nil {
				// Calculate connection duration and idle time for debugging
				connDuration := time.Since(w.connStartTime)
				timeSinceLastEvent := time.Since(lastHeartbeat)

				// Get ping/pong stats
				w.pongMutex.Lock()
				pingsTotal := w.pingsSent
				pongsTotal := w.pongsReceived
				timeSinceLastPong := time.Since(w.lastPongTime)
				w.pongMutex.Unlock()

				// Calculate ping/pong success rate
				var pongRate float64
				if pingsTotal > 0 {
					pongRate = float64(pongsTotal) / float64(pingsTotal) * 100
				}

				// Determine diagnosis
				var diagnosis string
				if pongRate >= 95 && timeSinceLastPong < 60*time.Second {
					diagnosis = "Connection was healthy (good ping/pong), likely server-side timeout or network interruption"
				} else if timeSinceLastPong > 60*time.Second {
					diagnosis = "Connection died (no pong for >60s), network issue detected"
				} else if pongRate < 80 {
					diagnosis = "Connection unstable (low pong rate), network quality issues"
				} else {
					diagnosis = "Connection closed unexpectedly"
				}

				// Log detailed context about the failure
				slog.Info("Jetstream connection closed", "duration", connDuration, "events_in_last_30s", eventCount, "time_since_last_event", timeSinceLastEvent, "pongs_received", pongsTotal, "pings_sent", pingsTotal, "pong_rate_pct", pongRate, "time_since_last_pong", timeSinceLastPong, "error", err, "diagnosis", diagnosis)

				return fmt.Errorf("failed to read message: %w", err)
			}

			// For now, process uncompressed messages
			// TODO: Re-enable compression once debugging is complete
			_ = decoder // Keep decoder to avoid unused variable error

			if err := w.processMessage(message); err != nil {
				slog.Error("ERROR processing message", "error", err)
				// Continue processing other messages
			} else {
				eventCount++
			}
		}
	}
}

// SetEventCallback sets a callback to be called for each event
func (w *Worker) SetEventCallback(cb EventCallback) {
	w.eventCallback = cb
}

// GetLastCursor returns the last processed cursor (time_us) for reconnects
func (w *Worker) GetLastCursor() int64 {
	w.cursorMutex.RLock()
	defer w.cursorMutex.RUnlock()
	return w.lastCursor
}

// processMessage processes a single Jetstream event
func (w *Worker) processMessage(message []byte) error {
	var event JetstreamEvent
	if err := json.Unmarshal(message, &event); err != nil {
		return fmt.Errorf("failed to unmarshal event: %w", err)
	}

	// Update cursor for reconnects (do this first, even if processing fails)
	w.cursorMutex.Lock()
	w.lastCursor = event.TimeUS
	w.cursorMutex.Unlock()

	// Call callback if set
	if w.eventCallback != nil {
		w.eventCallback(event.TimeUS)
	}

	// Process based on event kind
	switch event.Kind {
	case "commit":
		commit := event.Commit
		if commit == nil {
			return nil
		}

		// Set DID on commit from parent event
		commit.DID = event.DID

		// Debug: log first few collections we see to understand what's coming through
		if w.debugCollectionCount < 5 {
			slog.Debug("Jetstream received collection", "collection", commit.Collection, "did", commit.DID)
			w.debugCollectionCount++
		}

		// Process based on collection
		switch commit.Collection {
		case atproto.ManifestCollection:
			slog.Info("Jetstream processing manifest event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey)
			return w.processManifest(commit)
		case atproto.TagCollection:
			slog.Info("Jetstream processing tag event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey)
			return w.processTag(commit)
		case atproto.StarCollection:
			slog.Info("Jetstream processing star event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey)
			return w.processStar(commit)
		case atproto.RepoPageCollection:
			slog.Info("Jetstream processing repo page event", "did", commit.DID, "operation", commit.Operation, "rkey", commit.RKey)
			return w.processRepoPage(commit)
		default:
			// Ignore other collections
			return nil
		}

	case "identity":
		if event.Identity == nil {
			return nil
		}
		return w.processIdentity(&event)

	case "account":
		if event.Account == nil {
			return nil
		}
		return w.processAccount(&event)

	default:
		// Ignore unknown event kinds
		return nil
	}
}

// processManifest processes a manifest commit event
func (w *Worker) processManifest(commit *CommitEvent) error {
	// Resolve and upsert user with handle/PDS endpoint
	if err := w.processor.EnsureUser(context.Background(), commit.DID); err != nil {
		return fmt.Errorf("failed to ensure user: %w", err)
	}

	if commit.Operation == "delete" {
		// Delete manifest - rkey is just the digest, repository is not encoded
		digest := commit.RKey
		if err := db.DeleteManifest(w.db, commit.DID, "", digest); err != nil {
			return err
		}
		// Clean up any orphaned tags pointing to this manifest
		return db.CleanupOrphanedTags(w.db, commit.DID)
	}

	// Parse manifest record
	if commit.Record == nil {
		return nil // No record data, can't process
	}

	// Marshal map to bytes for processing
	recordBytes, err := json.Marshal(commit.Record)
	if err != nil {
		return fmt.Errorf("failed to marshal record: %w", err)
	}

	// Use shared processor for DB operations
	_, err = w.processor.ProcessManifest(context.Background(), commit.DID, recordBytes)
	return err
}

// processTag processes a tag commit event
func (w *Worker) processTag(commit *CommitEvent) error {
	// Resolve and upsert user with handle/PDS endpoint
	if err := w.processor.EnsureUser(context.Background(), commit.DID); err != nil {
		return fmt.Errorf("failed to ensure user: %w", err)
	}

	if commit.Operation == "delete" {
		// Delete tag - decode rkey back to repository and tag
		repo, tag := atproto.RKeyToRepositoryTag(commit.RKey)
		slog.Info("Jetstream deleting tag", "did", commit.DID, "repository", repo, "tag", tag, "rkey", commit.RKey)
		if err := db.DeleteTag(w.db, commit.DID, repo, tag); err != nil {
			slog.Error("Jetstream ERROR deleting tag", "error", err)
			return err
		}
		slog.Info("Jetstream successfully deleted tag", "did", commit.DID, "repository", repo, "tag", tag)
		return nil
	}

	// Parse tag record
	if commit.Record == nil {
		return nil
	}

	// Marshal map to bytes for processing
	recordBytes, err := json.Marshal(commit.Record)
	if err != nil {
		return fmt.Errorf("failed to marshal record: %w", err)
	}

	// Use shared processor for DB operations
	return w.processor.ProcessTag(context.Background(), commit.DID, recordBytes)
}

// processStar processes a star commit event
func (w *Worker) processStar(commit *CommitEvent) error {
	// Resolve and upsert the user who starred (starrer)
	if err := w.processor.EnsureUser(context.Background(), commit.DID); err != nil {
		return fmt.Errorf("failed to ensure user: %w", err)
	}

	if commit.Operation == "delete" {
		// Unstar - parse the rkey to get the subject (owner DID and repository)
		// Delete events don't include the full record, but the rkey contains the info we need
		ownerDID, repository, err := atproto.ParseStarRecordKey(commit.RKey)
		if err != nil {
			return fmt.Errorf("failed to parse star rkey: %w", err)
		}

		// Delete the star record
		return db.DeleteStar(w.db, commit.DID, ownerDID, repository)
	}

	// Parse star record
	if commit.Record == nil {
		return nil
	}

	// Marshal map to bytes for processing
	recordBytes, err := json.Marshal(commit.Record)
	if err != nil {
		return fmt.Errorf("failed to marshal record: %w", err)
	}

	// Use shared processor for DB operations
	return w.processor.ProcessStar(context.Background(), commit.DID, recordBytes)
}

// processRepoPage processes a repo page commit event
func (w *Worker) processRepoPage(commit *CommitEvent) error {
	// Resolve and upsert user with handle/PDS endpoint
	if err := w.processor.EnsureUser(context.Background(), commit.DID); err != nil {
		return fmt.Errorf("failed to ensure user: %w", err)
	}

	isDelete := commit.Operation == "delete"

	if isDelete {
		// Delete - rkey is the repository name
		slog.Info("Jetstream deleting repo page", "did", commit.DID, "repository", commit.RKey)
		if err := w.processor.ProcessRepoPage(context.Background(), commit.DID, commit.RKey, nil, true); err != nil {
			slog.Error("Jetstream ERROR deleting repo page", "error", err)
			return err
		}
		slog.Info("Jetstream successfully deleted repo page", "did", commit.DID, "repository", commit.RKey)
		return nil
	}

	// Parse repo page record
	if commit.Record == nil {
		return nil
	}

	// Marshal map to bytes for processing
	recordBytes, err := json.Marshal(commit.Record)
	if err != nil {
		return fmt.Errorf("failed to marshal record: %w", err)
	}

	// Use shared processor for DB operations
	return w.processor.ProcessRepoPage(context.Background(), commit.DID, commit.RKey, recordBytes, false)
}

// processIdentity processes an identity event (handle change)
func (w *Worker) processIdentity(event *JetstreamEvent) error {
	if event.Identity == nil {
		return nil
	}

	identity := event.Identity
	// Process via shared processor (only ATCR users will be logged at Info level)
	return w.processor.ProcessIdentity(context.Background(), identity.DID, identity.Handle)
}

// processAccount processes an account event (status change)
func (w *Worker) processAccount(event *JetstreamEvent) error {
	if event.Account == nil {
		return nil
	}

	account := event.Account
	// Process via shared processor (only ATCR users will be logged at Info level)
	return w.processor.ProcessAccount(context.Background(), account.DID, account.Active, account.Status)
}

// JetstreamEvent represents a Jetstream event
type JetstreamEvent struct {
	DID      string        `json:"did"`
	TimeUS   int64         `json:"time_us"`
	Kind     string        `json:"kind"` // "commit", "identity", "account"
	Commit   *CommitEvent  `json:"commit,omitempty"`
	Identity *IdentityInfo `json:"identity,omitempty"`
	Account  *AccountInfo  `json:"account,omitempty"`
}

// CommitEvent represents a commit event (create/update/delete)
type CommitEvent struct {
	Rev        string         `json:"rev"`
	Operation  string         `json:"operation"` // "create", "update", "delete"
	Collection string         `json:"collection"`
	RKey       string         `json:"rkey"`
	Record     map[string]any `json:"record,omitempty"`
	CID        string         `json:"cid,omitempty"`
	DID        string         `json:"-"` // Set from parent event
}

// IdentityInfo represents an identity event
type IdentityInfo struct {
	DID    string `json:"did"`
	Handle string `json:"handle"`
	Seq    int64  `json:"seq"`
	Time   string `json:"time"`
}

// AccountInfo represents an account status event
type AccountInfo struct {
	Active bool   `json:"active"`
	DID    string `json:"did"`
	Seq    int64  `json:"seq"`
	Time   string `json:"time"`
	Status string `json:"status,omitempty"`
}