cursor/cursor.go

package cursor

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/go-jet/jet/v2/mysql"
)

var (
	// ErrBadCursorString is returned when a cursor string cannot be decoded.
	ErrBadCursorString = errors.New("bad cursor string")

	columnRegistry     = make(map[ColumnKey]mysql.Column)
	columnRegistryLock sync.RWMutex
)

// CursorOrderCol is an interface for SQL expression that support the ORDER BY clause.
type CursorOrderCol interface {
	mysql.Column
	ASC() mysql.OrderByClause
	DESC() mysql.OrderByClause
}

// ColumnKey identifies a column and table in the database with a given value.
type ColumnKey struct {
	Table  string `json:"table"`
	Column string `json:"column"`
}

// NewColumnKey creates a new ColumnKey with the specified table and column names.
func NewColumnKey(col mysql.Column) ColumnKey {
	return ColumnKey{
		Table:  col.TableName(),
		Column: col.Name(),
	}
}

// IsEmpty returns true if the ColumnKey is empty.
func (k *ColumnKey) IsEmpty() bool {
	return k == nil || k.Table == "" || k.Column == ""
}

// String returns a string representation of the ColumnKey.
func (k ColumnKey) String() string {
	return fmt.Sprintf("%s.%s", k.Table, k.Column)
}

// RegisterColumn registers one or more columns with the cursor registry.
//
// Required at application startup before encoding, decoding, or SQL generation.
// Every index column, order column, and OrderValue column must be registered.
func RegisterColumn(columns ...mysql.Column) {
	columnRegistryLock.Lock()
	defer columnRegistryLock.Unlock()
	for _, col := range columns {
		columnRegistry[NewColumnKey(col)] = col
	}
}

// RegisterColumnList registers a list of columns with the cursor registry.
// See [RegisterColumn] for setup requirements.
func RegisterColumnList(columns mysql.ColumnList) {
	RegisterColumn([]mysql.Column(columns)...)
}

// GetColumn retrieves a column from the cursor registry.
func GetColumn(table, column string) (mysql.Column, error) {
	columnRegistryLock.RLock()
	defer columnRegistryLock.RUnlock()

	key := ColumnKey{
		Table:  table,
		Column: column,
	}

	col, ok := columnRegistry[key]
	if !ok {
		return nil, fmt.Errorf("column %s.%s not registered", table, column)
	}

	return col, nil
}

// GetColumnByKey retrieves a column from the cursor registry by its ColumnKey.
func GetColumnByKey(key ColumnKey) (mysql.Column, error) {
	columnRegistryLock.RLock()
	defer columnRegistryLock.RUnlock()

	col, ok := columnRegistry[key]
	if !ok {
		return nil, fmt.Errorf("column %s.%s not registered", key.Table, key.Column)
	}

	return col, nil
}

// OrderDirection represents the order direction of a cursor.
type OrderDirection int

const (
	// OrderAscending sorts results in ascending order.
	OrderAscending OrderDirection = iota
	// OrderDescending sorts results in descending order.
	OrderDescending
)

// MarshalJSON marshals the OrderDirection to JSON.
func (od OrderDirection) MarshalJSON() ([]byte, error) {
	switch od {
	case OrderAscending:
		return []byte(`"ASC"`), nil
	case OrderDescending:
		return []byte(`"DESC"`), nil
	default:
		return nil, fmt.Errorf("invalid order direction: %d", od)
	}
}

// UnmarshalJSON unmarshals the OrderDirection from JSON.
func (od *OrderDirection) UnmarshalJSON(data []byte) error {
	var dir string
	if err := json.Unmarshal(data, &dir); err != nil {
		return err
	}

	switch strings.ToUpper(dir) {
	case "ASC":
		*od = OrderAscending
	case "DESC":
		*od = OrderDescending
	default:
		return fmt.Errorf("invalid order direction: %s", dir)
	}
	return nil
}

// Compile-time checks to ensure value types implement ExprMarshaler
var _ ExprMarshaler[mysql.StringExpression, mysql.ColumnString] = (*StringValue)(nil)
var _ ExprMarshaler[mysql.IntegerExpression, mysql.ColumnInteger] = (*Int64Value)(nil)
var _ ExprMarshaler[mysql.IntegerExpression, mysql.ColumnInteger] = (*Uint64Value)(nil)
var _ ExprMarshaler[mysql.TimestampExpression, mysql.ColumnTimestamp] = (*TimestampValue)(nil)

// ExprMarshaler is the typed index-value extension point. Implement GenericExpr plus
// Expr and Col to produce go-jet expressions. Col panics when the column is not
// registered via [RegisterColumn].
type ExprMarshaler[E mysql.Expression, C mysql.Column] interface {
	GenericExpr
	Expr() E
	Col() C
}

// GenericExpr is the cursor value extension point shared by index and order columns.
// Implement ColumnKey, IsEmpty, and driver.Valuer. Order values decoded from cursor
// JSON must match a registered go-jet column kind; see package documentation.
type GenericExpr interface {
	ColumnKey() ColumnKey
	IsEmpty() bool
	driver.Valuer
}

var _ GenericExpr = (*StringValue)(nil)
var _ GenericExpr = (*Int64Value)(nil)
var _ GenericExpr = (*Uint64Value)(nil)
var _ GenericExpr = (*TimestampValue)(nil)

type StringValue struct {
	Key ColumnKey `json:"key"`
	Val string    `json:"val"`
}

// NewStringValue creates a cursor value for a string column (including UUIDs).
// The column must be registered via [RegisterColumn] before Col or decode use it.
func NewStringValue(val string, col mysql.ColumnString) *StringValue {
	return &StringValue{
		Key: NewColumnKey(col),
		Val: val,
	}
}

func (s StringValue) Expr() mysql.StringExpression {
	return mysql.String(s.Val)
}
func (s StringValue) Col() mysql.ColumnString {
	col, err := GetColumnByKey(s.Key)
	if err != nil {
		panic(err)
	}
	colStr, ok := col.(mysql.ColumnString)
	if !ok {
		panic(fmt.Errorf("column %s.%s is not a string column", s.Key.Table, s.Key.Column))
	}
	return colStr
}
func (s StringValue) Value() (driver.Value, error) {
	return s.Val, nil
}
func (s *StringValue) ColumnKey() ColumnKey {
	return s.Key
}
func (s *StringValue) IsEmpty() bool {
	return s == nil || s.Key.Table == "" || s.Key.Column == "" || s.Val == ""
}

type Int64Value struct {
	Key ColumnKey `json:"key"`
	Val int64     `json:"val"`
}

// NewInt64Value creates a cursor value for a signed integer column.
// The column must be registered via [RegisterColumn] before Col or decode use it.
func NewInt64Value(val int64, col mysql.ColumnInteger) *Int64Value {
	return &Int64Value{
		Key: NewColumnKey(col),
		Val: val,
	}
}

func (i Int64Value) Expr() mysql.IntegerExpression {
	return mysql.Int64(i.Val)
}
func (i Int64Value) Col() mysql.ColumnInteger {
	col, err := GetColumnByKey(i.Key)
	if err != nil {
		panic(err)
	}
	colInt, ok := col.(mysql.ColumnInteger)
	if !ok {
		panic(fmt.Errorf("column %s.%s is not an integer column", i.Key.Table, i.Key.Column))
	}
	return colInt
}
func (i Int64Value) Value() (driver.Value, error) {
	return i.Val, nil
}
func (i *Int64Value) ColumnKey() ColumnKey {
	return i.Key
}
func (i *Int64Value) IsEmpty() bool {
	return i == nil || i.Key.Table == "" || i.Key.Column == ""
}

type Uint64Value struct {
	Key ColumnKey `json:"key"`
	Val uint64    `json:"val"`
}

// NewUint64Value creates a cursor value for an unsigned integer column.
// The column must be registered via [RegisterColumn] before Col or decode use it.
func NewUint64Value(val uint64, col mysql.ColumnInteger) *Uint64Value {
	return &Uint64Value{
		Key: NewColumnKey(col),
		Val: val,
	}
}

func (i Uint64Value) Expr() mysql.IntegerExpression {
	return mysql.Uint64(i.Val)
}
func (i Uint64Value) Col() mysql.ColumnInteger {
	col, err := GetColumnByKey(i.Key)
	if err != nil {
		panic(err)
	}
	colInt, ok := col.(mysql.ColumnInteger)
	if !ok {
		panic(fmt.Errorf("column %s.%s is not an integer column", i.Key.Table, i.Key.Column))
	}
	return colInt
}
func (i Uint64Value) Value() (driver.Value, error) {
	return i.Val, nil
}
func (i *Uint64Value) ColumnKey() ColumnKey {
	return i.Key
}
func (i *Uint64Value) IsEmpty() bool {
	return i == nil || i.Key.Table == "" || i.Key.Column == ""
}

// TimestampValue stores a timestamp cursor position.
type TimestampValue struct {
	Key ColumnKey `json:"key"`
	Val time.Time `json:"val"`
}

// NewTimestampValue creates a timestamp cursor value for the given column.
// The column must be registered via [RegisterColumn] before Col or decode use it.
func NewTimestampValue(val time.Time, col mysql.ColumnTimestamp) *TimestampValue {
	return &TimestampValue{
		Key: NewColumnKey(col),
		Val: val,
	}
}

func (t TimestampValue) Expr() mysql.TimestampExpression {
	return mysql.TimestampT(t.Val)
}

func (t TimestampValue) Col() mysql.ColumnTimestamp {
	col, err := GetColumnByKey(t.Key)
	if err != nil {
		panic(err)
	}
	colTS, ok := col.(mysql.ColumnTimestamp)
	if !ok {
		panic(fmt.Errorf("column %s.%s is not a timestamp column", t.Key.Table, t.Key.Column))
	}
	return colTS
}

func (t TimestampValue) Value() (driver.Value, error) {
	return t.Val, nil
}

func (t *TimestampValue) ColumnKey() ColumnKey {
	return t.Key
}

func (t *TimestampValue) IsEmpty() bool {
	return t == nil || t.Key.Table == "" || t.Key.Column == "" || t.Val.IsZero()
}

// GenericCursor is the cursor extension point used by edge building, page counting,
// and pagination helpers without concrete type parameters. [Cursor] implements this
// interface.
type GenericCursor interface {
	IsEmpty() bool
	IsComposite() bool
	UsesTupleOrdering() bool
	OrderCol() CursorOrderCol
	Encode() (string, error)
	Decode(src string) error
	String() string
	GenericIndex() GenericExpr
	GenericOrderValue() GenericExpr
	Direction() OrderDirection
}

// Compile-time check to ensure that Cursor implements the GenericCursor interface.
var _ GenericCursor = (*Cursor[mysql.Expression, mysql.Column])(nil)

// Cursor identifies a location and order in the database.
//
// Index is the stable position column (typically the primary key). OrderColumnKey
// is the primary sort column shown to users. When they differ and OrderValue is
// set, pagination uses lexicographic tuple comparison on (order_col, index_col).
type Cursor[IE mysql.Expression, IC mysql.Column] struct {
	// Index identifies the column used for positioning the cursor and its current value.
	Index ExprMarshaler[IE, IC] `json:"index"`

	// OrderValue holds the order column value at this cursor position when using
	// composite tuple pagination.
	OrderValue GenericExpr `json:"order_val,omitempty"`

	// OrderColumnKey identifies the column used for ordering the results.
	OrderColumnKey ColumnKey `json:"order_col"`
	// OrderDir is the direction of the order (ASC or DESC).
	OrderDir OrderDirection `json:"order_dir"`
}

type cursorJSON struct {
	Index          json.RawMessage `json:"index"`
	OrderValue     json.RawMessage `json:"order_val"`
	OrderColumnKey ColumnKey       `json:"order_col"`
	OrderDir       OrderDirection  `json:"order_dir"`
}

// NewCursor creates a new Cursor with the specified parameters.
func NewCursor[IE mysql.Expression, IC mysql.Column](index ExprMarshaler[IE, IC], orderCol CursorOrderCol, orderDir OrderDirection) *Cursor[IE, IC] {
	return &Cursor[IE, IC]{
		Index:          index,
		OrderColumnKey: NewColumnKey(orderCol),
		OrderDir:       orderDir,
	}
}

// NewCursorFromAfterPtr decodes a Relay-style after cursor. Nil or empty after
// returns (nil, nil). Returns [ErrBadCursorString] when after is not valid cursor JSON.
func NewCursorFromAfterPtr[IE mysql.Expression, IC mysql.Column](
	newZero func() *Cursor[IE, IC],
	after *string,
) (*Cursor[IE, IC], error) {
	if after == nil || *after == "" {
		return nil, nil
	}
	c := newZero()
	return c, c.Decode(*after)
}

// NewCursorFromJSON returns a Cursor from a JSON representation.
// Returns [ErrBadCursorString] when src is not valid cursor JSON.
func NewCursorFromJSON[IE mysql.Expression, IC mysql.Column](zeroIndex ExprMarshaler[IE, IC], src []byte) (*Cursor[IE, IC], error) {
	cursor := Cursor[IE, IC]{
		Index: zeroIndex,
	}

	if err := decodeCursorJSON(&cursor, src, false, OrderAscending); err != nil {
		return nil, err
	}
	return &cursor, nil
}

// CopyWithVal returns a new cursor with the specified index value and the current
// ordering. Use for simple pagination where the index and order column are the same.
func (c *Cursor[IE, IC]) CopyWithVal(val ExprMarshaler[IE, IC]) *Cursor[IE, IC] {
	return NewCursor(val, c.OrderCol(), c.OrderDir)
}

// CopyWithVals returns a new cursor with the specified index and order values.
// Use when building edge cursors under tuple ordering: set both values so
// [IsComposite] is true and after pagination compares (order_col, index_col).
func (c *Cursor[IE, IC]) CopyWithVals(index ExprMarshaler[IE, IC], orderVal GenericExpr) *Cursor[IE, IC] {
	result := NewCursor(index, c.OrderCol(), c.OrderDir)
	result.OrderValue = orderVal
	return result
}

// UsesTupleOrdering reports whether results are sorted by (order_col, index_col).
//
// True when the order column differs from the index column. Applies to default
// cursors on the first page and drives [OrderByClauses]. Unlike [IsComposite],
// OrderValue is not required.
func (c *Cursor[IE, IC]) UsesTupleOrdering() bool {
	if c == nil || c.Index == nil {
		return false
	}
	indexKey := c.Index.ColumnKey()
	if indexKey.IsEmpty() || c.OrderColumnKey.IsEmpty() {
		return false
	}
	return c.OrderColumnKey != indexKey
}

// IsComposite reports whether after-pagination filters on (order_col, index_col).
//
// True when tuple ordering is active and OrderValue is set—typical for encoded edge
// cursors when the sort column is non-unique. Drives [PaginateConds] and page
// counting via lexicographic tuple comparison.
func (c *Cursor[IE, IC]) IsComposite() bool {
	if c == nil || c.Index == nil || c.Index.IsEmpty() {
		return false
	}
	if c.OrderColumnKey == c.Index.ColumnKey() {
		return false
	}
	return c.OrderValue != nil && !c.OrderValue.IsEmpty()
}

// IsEmpty returns true if the cursor or any of its keys are empty or unmapped.
func (c *Cursor[IE, IC]) IsEmpty() bool {
	if c == nil || c.Index == nil || c.Index.IsEmpty() {
		return true
	}
	_, err := GetColumnByKey(c.OrderColumnKey)
	return err != nil
}

// GenericIndex returns the Index ExprMarshaler as a GenericExpr.
func (c *Cursor[IE, IC]) GenericIndex() GenericExpr {
	return c.Index
}

// GenericOrderValue returns the order column value when using composite pagination.
func (c *Cursor[IE, IC]) GenericOrderValue() GenericExpr {
	return c.OrderValue
}

// Direction returns the order direction expected by the cursor.
func (c *Cursor[IE, IC]) Direction() OrderDirection {
	return c.OrderDir
}

// OrderCol returns the column used for ordering the results. Panics if the
// column is not registered.
func (c *Cursor[IE, IC]) OrderCol() CursorOrderCol {
	col, err := GetColumnByKey(c.OrderColumnKey)
	if err != nil {
		panic(err)
	}
	return col
}

// Encode returns a stringified JSON representation of the cursor.
func (c *Cursor[IE, IC]) Encode() (string, error) {
	bytes, err := json.Marshal(c)
	if err != nil {
		return "", fmt.Errorf("failed to marshal cursor: %w", err)
	}
	return string(bytes), nil
}

// Decode decodes a stringified JSON representation of the cursor into this object.
// Returns [ErrBadCursorString] when the input is not valid cursor JSON.
func (c *Cursor[IE, IC]) Decode(src string) error {
	if c == nil {
		return fmt.Errorf("cursor is nil")
	}
	return decodeCursorJSON(c, []byte(src), false, OrderAscending)
}

// DecodeAndOrder decodes a cursor string like [Cursor.Decode] but replaces the
// encoded order direction with orderDir.
func (c *Cursor[IE, IC]) DecodeAndOrder(src string, orderDir OrderDirection) error {
	return decodeCursorJSON(c, []byte(src), true, orderDir)
}

func decodeCursorJSON[IE mysql.Expression, IC mysql.Column](
	c *Cursor[IE, IC],
	src []byte,
	overrideDir bool,
	orderDir OrderDirection,
) error {
	var raw cursorJSON
	if err := json.Unmarshal(src, &raw); err != nil {
		return fmt.Errorf("%w: invalid JSON: %v", ErrBadCursorString, err)
	}

	if err := json.Unmarshal(raw.Index, c.Index); err != nil {
		return fmt.Errorf("%w: invalid index: %v", ErrBadCursorString, err)
	}

	c.OrderColumnKey = raw.OrderColumnKey
	c.OrderDir = raw.OrderDir
	if overrideDir {
		c.OrderDir = orderDir
	}

	if len(raw.OrderValue) > 0 {
		orderVal, err := unmarshalGenericExpr(raw.OrderValue, raw.OrderColumnKey)
		if err != nil {
			return fmt.Errorf("%w: invalid order value: %v", ErrBadCursorString, err)
		}
		c.OrderValue = orderVal
	} else {
		c.OrderValue = nil
	}

	return nil
}

func unmarshalGenericExpr(data []byte, key ColumnKey) (GenericExpr, error) {
	col, err := GetColumnByKey(key)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrBadCursorString, err)
	}

	switch col.(type) {
	case mysql.ColumnTimestamp:
		var v TimestampValue
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, fmt.Errorf("%w: %v", ErrBadCursorString, err)
		}
		return &v, nil
	case mysql.ColumnInteger:
		var v Uint64Value
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, fmt.Errorf("%w: %v", ErrBadCursorString, err)
		}
		return &v, nil
	case mysql.ColumnString:
		var v StringValue
		if err := json.Unmarshal(data, &v); err != nil {
			return nil, fmt.Errorf("%w: %v", ErrBadCursorString, err)
		}
		return &v, nil
	default:
		return nil, fmt.Errorf("%w: unsupported order column type for %s", ErrBadCursorString, key)
	}
}

// String returns a string representation of the Cursor.
func (c *Cursor[IE, IC]) String() string {
	bytes, err := json.Marshal(c)
	if err != nil {
		return ""
	}
	return string(bytes)
}