package jetstream

import (
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strconv"
	"time"

	comatproto "github.com/bluesky-social/indigo/api/atproto"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
)

// bsky jetstream compresses with this dictionary, and hydrant's frames decode with it too
//
//go:embed zstd_dictionary
var zstdDictionary []byte

const (
	EventKindCommit   = "commit"
	EventKindAccount  = "account"
	EventKindIdentity = "identity"

	CommitOperationCreate = "create"
	CommitOperationUpdate = "update"
	CommitOperationDelete = "delete"
)

// hydrant pings every 30s, so 90s of silence means the connection is dead
const defaultReadTimeout = 90 * time.Second

// refused dials fail instantly, without the wait this loop spins hot
const reconnectDelay = time.Second

// indigo's jetstream event plus a live flag on commits, which is what hydrant's /subscribe sends
type Event struct {
	Did      string                                  `json:"did"`
	TimeUS   int64                                   `json:"time_us"`
	Kind     string                                  `json:"kind,omitempty"`
	Commit   *Commit                                 `json:"commit,omitempty"`
	Account  *comatproto.SyncSubscribeRepos_Account  `json:"account,omitempty"`
	Identity *comatproto.SyncSubscribeRepos_Identity `json:"identity,omitempty"`
}

type Commit struct {
	Rev        string          `json:"rev,omitempty"`
	Operation  string          `json:"operation,omitempty"`
	Collection string          `json:"collection,omitempty"`
	RKey       string          `json:"rkey,omitempty"`
	Record     json.RawMessage `json:"record,omitempty"`
	CID        string          `json:"cid,omitempty"`
	// plain jetstream never sends this, so nil counts as live
	Live *bool `json:"live,omitempty"`
}

func (c *Commit) IsLive() bool {
	return c.Live == nil || *c.Live
}

type StreamConfig struct {
	Endpoint    string
	Collections []string
	// hydrant only backfills when asked here, plain jetstream doesn't know the param
	EventTypes []string
	Compress   bool

	// read on every connect, ok false replays from 0 unless TailWhenNoCursor is set
	Cursor           func() (timeUs int64, ok bool, err error)
	TailWhenNoCursor bool
	// each hydrant stamps time_us its own way, so resume a little before the cursor
	Rewind time.Duration

	ReadTimeout    time.Duration
	ReconnectDelay time.Duration
	UserAgent      string
}

// the handler saves the cursor in the same transaction as the event, so a crash can't keep one without the other
type Stream struct {
	cfg StreamConfig
	l   *slog.Logger
	dec *zstd.Decoder
}

func NewStream(cfg StreamConfig, logger *slog.Logger) (*Stream, error) {
	if cfg.Cursor == nil {
		return nil, fmt.Errorf("stream needs a cursor source")
	}
	if cfg.ReadTimeout == 0 {
		cfg.ReadTimeout = defaultReadTimeout
	}
	if cfg.ReconnectDelay == 0 {
		cfg.ReconnectDelay = reconnectDelay
	}
	s := &Stream{cfg: cfg, l: logger}
	if cfg.Compress {
		dec, err := zstd.NewReader(nil, zstd.WithDecoderDicts(zstdDictionary))
		if err != nil {
			return nil, fmt.Errorf("creating zstd decoder: %w", err)
		}
		s.dec = dec
	}
	return s, nil
}

// a handler error drops the connection, so the event comes back from the cursor
func (s *Stream) Run(ctx context.Context, handle func(context.Context, *Event) error) error {
	for {
		if err := s.readOnce(ctx, handle); err != nil && ctx.Err() == nil {
			s.l.Error("reading stream", "err", err)
		}
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(s.cfg.ReconnectDelay):
		}
	}
}

// nil leaves the cursor param off, which starts at the live tail
func (s *Stream) resumeCursor() (*int64, error) {
	us, ok, err := s.cfg.Cursor()
	if err != nil {
		return nil, err
	}
	if !ok {
		if s.cfg.TailWhenNoCursor {
			return nil, nil
		}
		us = 0
	}
	cursor := max(us-s.cfg.Rewind.Microseconds(), 0)
	return &cursor, nil
}

func (s *Stream) url(cursor *int64) (string, error) {
	u, err := url.Parse(s.cfg.Endpoint)
	if err != nil {
		return "", fmt.Errorf("parsing endpoint %q: %w", s.cfg.Endpoint, err)
	}
	q := u.Query()
	q.Del("cursor")
	if cursor != nil {
		q.Set("cursor", strconv.FormatInt(*cursor, 10))
	}
	q.Del("wantedCollections")
	for _, nsid := range s.cfg.Collections {
		q.Add("wantedCollections", nsid)
	}
	q.Del("wantedEventTypes")
	for _, t := range s.cfg.EventTypes {
		q.Add("wantedEventTypes", t)
	}
	u.RawQuery = q.Encode()
	return u.String(), nil
}

func (s *Stream) readOnce(ctx context.Context, handle func(context.Context, *Event) error) error {
	cursor, err := s.resumeCursor()
	if err != nil {
		return fmt.Errorf("reading cursor: %w", err)
	}
	target, err := s.url(cursor)
	if err != nil {
		return err
	}

	header := http.Header{}
	if s.cfg.UserAgent != "" {
		header.Set("User-Agent", s.cfg.UserAgent)
	}
	if s.dec != nil {
		header.Set("Socket-Encoding", "zstd")
	}

	s.l.Info("connecting to stream", "url", target)
	con, _, err := websocket.DefaultDialer.DialContext(ctx, target, header)
	if err != nil {
		return err
	}
	defer con.Close()
	// a quiet stream blocks the read and never checks ctx, so cancel closes the conn
	defer context.AfterFunc(ctx, func() { con.Close() })()

	alive := func() error { return con.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeout)) }
	con.SetPingHandler(func(data string) error {
		if err := alive(); err != nil {
			return err
		}
		return con.WriteControl(websocket.PongMessage, []byte(data), time.Now().Add(time.Second))
	})

	for {
		if err := alive(); err != nil {
			return err
		}
		_, msg, err := con.ReadMessage()
		if err != nil {
			return fmt.Errorf("reading message: %w", err)
		}
		if s.dec != nil {
			if msg, err = s.dec.DecodeAll(msg, nil); err != nil {
				return fmt.Errorf("decompressing message: %w", err)
			}
		}
		var e Event
		if err := json.Unmarshal(msg, &e); err != nil {
			s.l.Warn("skipping undecodable frame", "err", err)
			continue
		}
		if err := handle(ctx, &e); err != nil {
			return fmt.Errorf("handling event at %d: %w", e.TimeUS, err)
		}
	}
}