package knotfeed

import (
	"cmp"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

const (
	livenessTimeout  = 90 * time.Second
	pongWriteTimeout = 5 * time.Second
	initialBackoff   = time.Second
	maxBackoff       = 60 * time.Second
	maxFrameBytes    = 8 << 20
	subscribeReposNS = "com.atproto.sync.subscribeRepos"
	eventsPath       = "/events"

	defaultMaxHandlerAttempts = 3
	defaultPoisonGrace        = 10 * time.Minute
)

type Feed struct {
	kind uint8
}

const (
	atprotoKind uint8 = iota
	legacyKind
)

var (
	FeedAtproto = Feed{kind: atprotoKind}
	FeedLegacy  = Feed{kind: legacyKind}
)

type feedSpec struct {
	token   string
	name    string
	path    string
	msgKind int
	other   Feed
	decode  func([]byte, *slog.Logger) (Message, error)
}

var feedSpecs = [...]feedSpec{
	atprotoKind: {
		token: "atproto", name: "atproto firehose", path: "/xrpc/" + subscribeReposNS,
		msgKind: websocket.BinaryMessage, other: FeedLegacy, decode: Decode,
	},
	legacyKind: {
		token: "legacy", name: "knot event stream", path: eventsPath,
		msgKind: websocket.TextMessage, other: FeedAtproto, decode: decodeLegacy,
	},
}

func (f Feed) spec() feedSpec { return feedSpecs[f.kind] }

func (f Feed) Token() string { return f.spec().token }

func (f Feed) String() string { return f.spec().name }

func (f Feed) other() Feed { return f.spec().other }

func (f Feed) Live(now time.Time) Cursor {
	if f == FeedLegacy {
		return NewCursor(f, now.UnixNano())
	}
	return NewCursor(f, 0)
}

var ErrUnrecognizedFeed = errors.New("unrecognized feed token")

func ParseFeed(token string) (Feed, error) {
	kind := slices.IndexFunc(feedSpecs[:], func(spec feedSpec) bool { return spec.token == token })
	if kind < 0 {
		return FeedAtproto, fmt.Errorf("%w: %q", ErrUnrecognizedFeed, token)
	}
	return Feed{kind: uint8(kind)}, nil
}

type Cursor struct {
	feed Feed
	seq  int64
}

func NewCursor(feed Feed, seq int64) Cursor {
	return Cursor{feed: feed, seq: seq}
}

func (c Cursor) Feed() Feed { return c.feed }

func (c Cursor) Seq() int64 { return c.seq }

func (c Cursor) Encode() string {
	return c.feed.Token() + ":" + strconv.FormatInt(c.seq, 10)
}

func ParseCursor(raw string) (Cursor, error) {
	token, rest, named := strings.Cut(raw, ":")
	if !named {
		token, rest = FeedAtproto.Token(), raw
	}
	feed, err := ParseFeed(token)
	if err != nil {
		return Cursor{}, err
	}
	seq, err := strconv.ParseInt(rest, 10, 64)
	if err != nil {
		return Cursor{}, fmt.Errorf("cursor %q: %w", raw, err)
	}
	return NewCursor(feed, seq), nil
}

type Consumer struct {
	Host   string
	NoTLS  bool
	Dialer *websocket.Dialer
	Logger *slog.Logger

	LoadCursor     func(context.Context) (Cursor, error)
	StoreCursor    func(context.Context, Cursor) error
	Handle         func(context.Context, Message) error
	OutdatedReplay func(context.Context, Feed) Cursor
	OnConnectError func(error)

	ReplayFromStart bool

	MaxHandlerAttempts int
	PoisonGrace        time.Duration

	now func() time.Time

	current        Feed
	adopted        bool
	proven         bool
	poisonSeq      int64
	poisonAttempts int
	poisonSince    time.Time
}

func (c *Consumer) Run(ctx context.Context) error {
	c.fill()
	backoff := initialBackoff
	for {
		if ctx.Err() != nil {
			return nil
		}
		advanced, resync, err := c.session(ctx)
		switch {
		case err != nil:
			c.reactTo(err)
		case resync:
			advanced = true
		}
		if advanced {
			backoff = initialBackoff
		} else {
			backoff = min(backoff*2, maxBackoff)
		}
		if resync {
			continue
		}
		timer := time.NewTimer(backoff)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil
		case <-timer.C:
		}
	}
}

func (c *Consumer) reactTo(err error) {
	var mismatch feedMismatch
	switch {
	case errors.As(err, &mismatch):
	case isConnectError(err):
		if c.OnConnectError != nil {
			c.OnConnectError(err)
		}
		if c.proven {
			c.Logger.Error("knot turned us away on a feed it had already served",
				"host", c.Host, "feed", c.current, "err", err)
			return
		}
	default:
		c.Logger.Error("session died for its own reasons",
			"host", c.Host, "feed", c.current, "err", err)
		return
	}
	c.Logger.Warn("trying the other feed, since this knot won't serve the one we asked for",
		"host", c.Host, "tried", c.current, "next", c.current.other(), "err", err)
	c.current, c.proven = c.current.other(), false
}

type connectError struct{ err error }

func (e connectError) Error() string { return e.err.Error() }

func (e connectError) Unwrap() error { return e.err }

func isConnectError(err error) bool {
	var ce connectError
	return errors.As(err, &ce)
}

type feedMismatch struct{ feed Feed }

func (e feedMismatch) Error() string {
	return fmt.Sprintf("%s sent the %s's frames", e.feed, e.feed.other())
}

func (c *Consumer) session(ctx context.Context) (advanced bool, resync bool, err error) {
	stored, err := c.LoadCursor(ctx)
	if err != nil {
		return false, false, fmt.Errorf("loading cursor: %w", err)
	}
	if !c.adopted {
		c.adopted = true
		if stored.Seq() > 0 {
			c.current = stored.Feed()
		}
	}
	cursor := c.sessionCursor(stored)
	conn, _, err := c.Dialer.DialContext(ctx, c.url(cursor), nil)
	if err != nil {
		return false, false, connectError{err}
	}
	c.proven = true
	defer conn.Close()
	watcher := context.AfterFunc(ctx, func() { conn.Close() })
	defer watcher()
	conn.SetReadLimit(maxFrameBytes)
	c.Logger.Info("subscribed to knot feed", "host", c.Host, "feed", c.current, "cursor", cursor)

	conn.SetPingHandler(func(payload string) error {
		conn.SetReadDeadline(time.Now().Add(livenessTimeout))
		return conn.WriteControl(websocket.PongMessage, []byte(payload), time.Now().Add(pongWriteTimeout))
	})
	conn.SetReadDeadline(time.Now().Add(livenessTimeout))
	want := c.current.spec().msgKind
	for {
		kind, data, err := conn.ReadMessage()
		if err != nil {
			return advanced, false, err
		}
		conn.SetReadDeadline(time.Now().Add(livenessTimeout))
		if kind != want {
			return advanced, false, feedMismatch{feed: c.current}
		}
		message, err := c.decode(data)
		if err != nil {
			c.Logger.Error("undecodable knot feed frame", "host", c.Host, "feed", c.current, "err", err)
			continue
		}
		resync, err := c.dispatch(ctx, message, &cursor)
		if resync {
			return advanced, true, nil
		}
		if err != nil {
			return false, false, err
		}
		if message.Type == TypeCommit {
			advanced = true
		}
	}
}

func (c *Consumer) decode(data []byte) (Message, error) {
	message, err := c.current.spec().decode(data, c.Logger)
	message.Feed = c.current
	return message, err
}

func (c *Consumer) dispatch(ctx context.Context, message Message, cursor *Cursor) (bool, error) {
	if message.Type == TypeCommit {
		return c.commit(ctx, message, cursor)
	}
	if c.unreachable(message) {
		c.Logger.Warn("firehose can't serve our cursor, resuming live", "host", c.Host, "error", c.resumeReason(message))
		*cursor = c.replayCursor(ctx)
		if err := c.store(ctx, *cursor); err != nil {
			return false, err
		}
		return true, nil
	}
	if c.Handle != nil {
		if err := c.Handle(ctx, message); err != nil {
			c.Logger.Error("firehose handler failed", "host", c.Host, "frame", message.Type, "err", err)
		}
	}
	return false, nil
}

func (c *Consumer) replayCursor(ctx context.Context) Cursor {
	live := c.current.Live(c.now())
	if c.OutdatedReplay == nil {
		return live
	}
	replay := c.OutdatedReplay(ctx, c.current)
	if replay.Feed() != c.current {
		c.Logger.Warn("replay hook named the other feed, so resuming live instead",
			"host", c.Host, "feed", c.current, "replay", replay.Feed())
		return live
	}
	return replay
}

func (c *Consumer) commit(ctx context.Context, message Message, cursor *Cursor) (bool, error) {
	if c.Handle != nil {
		if err := c.Handle(ctx, message); err != nil {
			if !c.poison(message.Commit.Seq) {
				return false, fmt.Errorf("frame seq %d: %w", message.Commit.Seq, err)
			}
			c.Logger.Error("handler keeps failing at one frame, skipping it",
				"host", c.Host, "seq", message.Commit.Seq, "attempts", c.poisonAttempts, "err", err)
		}
	}
	*cursor = NewCursor(c.current, message.Commit.Seq)
	if err := c.store(ctx, *cursor); err != nil {
		return false, err
	}
	return false, nil
}

func (c *Consumer) store(ctx context.Context, cursor Cursor) error {
	if err := c.StoreCursor(ctx, cursor); err != nil {
		return fmt.Errorf("storing cursor: %w", err)
	}
	return nil
}

func (c *Consumer) poison(seq int64) bool {
	if seq != c.poisonSeq {
		c.poisonSeq = seq
		c.poisonAttempts = 0
		c.poisonSince = c.now()
	}
	c.poisonAttempts++
	return c.poisonAttempts >= c.MaxHandlerAttempts && c.now().Sub(c.poisonSince) >= c.PoisonGrace
}

func (c *Consumer) unreachable(message Message) bool {
	if message.Type == TypeError {
		return message.Error == "FutureCursor"
	}
	return message.Type == TypeInfo && message.InfoName == "OutdatedCursor"
}

func (c *Consumer) resumeReason(message Message) string {
	if message.Type == TypeError {
		return message.Error
	}
	return message.InfoName
}

func (c *Consumer) url(cursor Cursor) string {
	scheme := "wss"
	if c.NoTLS {
		scheme = "ws"
	}
	endpoint := url.URL{Scheme: scheme, Host: c.Host, Path: c.current.spec().path}
	if seq := cursor.Seq(); seq > 0 || (seq == 0 && c.ReplayFromStart) {
		endpoint.RawQuery = "cursor=" + strconv.FormatInt(seq, 10)
	}
	return endpoint.String()
}

func (c *Consumer) sessionCursor(stored Cursor) Cursor {
	switch {
	case stored.Feed() == c.current && stored.Seq() > 0:
		return stored
	case c.ReplayFromStart:
		return NewCursor(c.current, 0)
	default:
		return c.current.Live(c.now())
	}
}

func (c *Consumer) fill() {
	c.MaxHandlerAttempts = cmp.Or(c.MaxHandlerAttempts, defaultMaxHandlerAttempts)
	c.PoisonGrace = cmp.Or(c.PoisonGrace, defaultPoisonGrace)
	c.Dialer = cmp.Or(c.Dialer, websocket.DefaultDialer)
	c.Logger = cmp.Or(c.Logger, slog.Default())
	if c.now == nil {
		c.now = time.Now
	}
}