api

type AggrFn int

type BindConfig func(set *pflag.FlagSet) error

type CollectNotification struct {
	FQN         string `json:"fqn"`
	EncodedKeys string `json:"encoded_keys"`
	Bucket      string `json:"bucket,omitempty"`
}

type CollectNotifierFactory NotifierFactory[CollectNotification]

type ContextKey int

type DataSource struct {
	FQN    string                 `json:"fqn"`
	Kind   string                 `json:"kind"`
	Config manifests.ParsedConfig `json:"config"`
}

type DataSourceGetter interface {
	GetDataSource(FQN string) (DataSource, error)
}

type DataSourceManager interface {
	BindDataSource(fd DataSource) error
	UnbindDataSource(FQN string) error
	HasDataSource(FQN string) bool
}

type DataSourceReconcile func(ctx context.Context, rr DataSourceReconcileRequest) (changed bool, err error)

type DataSourceReconcileRequest struct {
	DataSource     *manifests.DataSource
	RuntimeManager RuntimeManager
	Client         client.Client
	Scheme         *runtime.Scheme
	CoreAddress    string
}

type Engine interface {
	// FeatureDescriptor returns the FeatureDescriptor for the given FQN
	FeatureDescriptor(ctx context.Context, selector string) (FeatureDescriptor, error)
	// Get returns the value for the given FQN and keys
	// If the feature is not available, it returns nil.
	// If the feature is windowed, the returned Value is a map from window function to Value.
	Get(ctx context.Context, selector string, keys Keys) (Value, FeatureDescriptor, error)
	// Set sets the raw value for the given FQN and keys
	// If the feature's primitive is a List, it replaces the entire list.
	// If the feature is windowed, it is aliased to WindowAdd instead of Set.
	Set(ctx context.Context, FQN string, keys Keys, val any, ts time.Time) error
	// Append appends to the raw value for the given FQN and keys
	// If the feature's primitive is NOT a List it will throw an error.
	Append(ctx context.Context, FQN string, keys Keys, val any, ts time.Time) error
	// Incr increments the raw value of the feature.
	// If the feature's primitive is NOT a Scalar it will throw an error.
	// It returns the updated value in the state, and an error if occurred.
	Incr(ctx context.Context, FQN string, keys Keys, by any, ts time.Time) error
	// Update is the common function to update a feature SimpleValue.
	// Under the hood, it utilizes lower-level functions depending on the type of the feature.
	//  - Set for Scalars
	//	- Append for Lists
	//  - WindowAdd for Windows
	Update(ctx context.Context, FQN string, keys Keys, val any, ts time.Time) error
}

type ExtendedManager interface {
	Engine
	RuntimeManager
	DataSourceGetter
}

type FeatureApply func(fd FeatureDescriptor, builder manifests.FeatureBuilder, pl Pipeliner, src ExtendedManager) error

type FeatureDescriptor struct {
	FQN          string        `json:"FQN"`
	Primitive    PrimitiveType `json:"primitive"`
	Aggr         []AggrFn      `json:"aggr"`
	Freshness    time.Duration `json:"freshness"`
	Staleness    time.Duration `json:"staleness"`
	Timeout      time.Duration `json:"timeout"`
	KeepPrevious *KeepPrevious `json:"keep_previous"`
	Keys         []string      `json:"keys"`
	Builder      string        `json:"builder"`
	RuntimeEnv   string        `json:"runtimeEnv"`
	DataSource   string        `json:"data_source"`
	Dependencies []string      `json:"dependencies"`
}

type FeatureDescriptorGetter func(ctx context.Context, FQN string) (FeatureDescriptor, error)

type FeatureManager interface {
	BindFeature(in *manifests.Feature) error
	UnbindFeature(FQN string) error
	HasFeature(FQN string) bool
}

type HistoricalWriter interface {
	Commit(context.Context, WriteNotification) error
	Flush(ctx context.Context, fqn string) error
	FlushAll(context.Context) error
	Close(ctx context.Context) error
	BindFeature(fd *FeatureDescriptor, model *manifests.ModelSpec, getter FeatureDescriptorGetter) error
}

type HistoricalWriterFactory func(viper *viper.Viper) (HistoricalWriter, error)

type KeepPrevious struct {
	Versions uint
	Over     time.Duration
}

type Keys map[string]string

type Logger interface {
	Logger() logr.Logger
}

type LowLevelValue interface {
	~int | ~string | ~float64 | time.Time | ~[]int | ~[]string | ~[]float64 | ~[]time.Time | WindowResultMap
}

type ManagerEngine interface {
	Logger
	FeatureManager
	DataSourceManager
	RuntimeManager
	Engine
}

type Middleware func(next MiddlewareHandler) MiddlewareHandler

type MiddlewareHandler func(ctx context.Context, fd FeatureDescriptor, keys Keys, val Value) (Value, error)

type ModelDescriptor struct {
	Features        []string               `json:"features"`
	KeyFeature      string                 `json:"keyFeature,omitempty"`
	Keys            []string               `json:"keys"`
	ModelFramework  string                 `json:"modelFramework"`
	ModelServer     string                 `json:"modelServer"`
	InferenceConfig manifests.ParsedConfig `json:"inferenceConfig"`
}

type ModelReconcileRequest struct {
	Model  *manifests.Model
	Client client.Client
	Scheme *runtime.Scheme
}

type ModelServer interface {
	Reconcile(ctx context.Context, rr ModelReconcileRequest) (changed bool, err error)
	Owns() []client.Object
	Serve(ctx context.Context, fd FeatureDescriptor, md ModelDescriptor, val Value) (Value, error)
}

type Notification interface {
	CollectNotification | WriteNotification
}

type Notifier[T Notification] interface {
	Notify(context.Context, T) error
	Subscribe(context.Context) (<-chan T, error)
}

type NotifierFactory[T Notification] func(viper *viper.Viper) (Notifier[T], error)

type ParsedProgram struct {
	// Primitive is the primitive that this program is returning
	Primitive PrimitiveType

	// Dependencies is a list of FQNs that this program *might* be depended on
	Dependencies []string
}

type Pipeliner interface {
	AddPreGetMiddleware(priority int, fn Middleware)
	AddPostGetMiddleware(priority int, fn Middleware)
	AddPreSetMiddleware(priority int, fn Middleware)
	AddPostSetMiddleware(priority int, fn Middleware)
}

type Plugins interface {
	BindConfig | FeatureApply | DataSourceReconcile | StateFactory |
		CollectNotifierFactory | WriteNotifierFactory |
		HistoricalWriterFactory
}

type PrimitiveType int

type RawBucket struct {
	FQN         string          `json:"FQN"`
	Bucket      string          `json:"bucket"`
	EncodedKeys string          `json:"encoded_keys"`
	Data        WindowResultMap `json:"raw"`
}

type RawBuckets []RawBucket

type RuntimeManager interface {
	// LoadProgram loads a program into the runtime.
	LoadProgram(env, fqn, program string, packages []string) (*ParsedProgram, error)

	// ExecuteProgram executes a program in the runtime.
	ExecuteProgram(ctx context.Context, env string, fqn string, keys Keys, row map[string]any, ts time.Time, dryRun bool) (value Value, keyz Keys, err error)

	// GetSidecars returns the sidecar containers attached to the current container.
	GetSidecars() []v1.Container

	// GetDefaultEnv returns the default environment for the current container.
	GetDefaultEnv() string
}

type State interface {
	// Get returns the SimpleValue of the feature.
	// If the feature is not available, it returns nil.
	// If the feature is windowed, the returned SimpleValue is a map from window function to SimpleValue.
	// version indicates the previous version of the feature. If version is 0, the latest version is returned.
	Get(ctx context.Context, fd FeatureDescriptor, keys Keys, version uint) (*Value, error)

	// Set sets the SimpleValue of the feature.
	// If the feature's primitive is a List, it replaces the entire list.
	// If the feature is windowed, it is aliased to WindowAdd instead of Set.
	Set(ctx context.Context, fd FeatureDescriptor, keys Keys, val any, timestamp time.Time) error

	// Append appends the SimpleValue to the feature.
	// If the feature's primitive is NOT a List it will throw an error.
	Append(ctx context.Context, fd FeatureDescriptor, keys Keys, val any, ts time.Time) error

	// Incr increments the SimpleValue of the feature.
	// If the feature's primitive is NOT a Scalar it will throw an error.
	// It returns the updated value in the state, and an error if occurred.
	Incr(ctx context.Context, fd FeatureDescriptor, keys Keys, by any, timestamp time.Time) error
	// Update is the common function to update a feature SimpleValue.
	// Under the hood, it utilizes lower-level functions depending on the type of the feature.
	//  - Set for Scalars
	//	- Append for Lists
	//  - WindowAdd for Windows
	Update(ctx context.Context, fd FeatureDescriptor, keys Keys, val any, timestamp time.Time) error

	// WindowAdd adds a Bucket to the window that contains aggregated data internally
	// Later the bucket's aggregations should be aggregated for the whole Window via Get
	//
	// Buckets should last *at least* as long as the feature's staleness time + DeadGracePeriod
	WindowAdd(ctx context.Context, fd FeatureDescriptor, keys Keys, val any, timestamp time.Time) error

	// WindowBuckets returns the list of RawBuckets for the feature and specific Keys.
	WindowBuckets(ctx context.Context, fd FeatureDescriptor, keys Keys, buckets []string) (RawBuckets, error)

	// DeadWindowBuckets returns the list of all the dead feature's RawBuckets of all the entities.
	DeadWindowBuckets(ctx context.Context, fd FeatureDescriptor, ignore RawBuckets) (RawBuckets, error)

	// Ping is a simple keepalive check for the state.
	// It should return an error in case an error occurred, or nil if everything is alright.
	Ping(ctx context.Context) error
}

type StateFactory func(viper *viper.Viper) (State, error)

type StateMethod int

type Value struct {
	// Value can be cast to LowLevelValue using the ToLowLevelValue() method
	Value     any       `json:"value"`
	Timestamp time.Time `json:"timestamp"`
	Fresh     bool      `json:"fresh"`
}

type WindowResultMap map[AggrFn]float64

type WriteNotification struct {
	FQN          string `json:"fqn"`
	EncodedKeys  string `json:"encoded_keys"`
	Bucket       string `json:"bucket,omitempty"`
	ActiveBucket bool   `json:"active_bucket,omitempty"`
	Value        *Value `json:"value,omitempty"`
}

type WriteNotifierFactory NotifierFactory[WriteNotification]