package operator

import (
	"context"
	"fmt"

	"flux/core"
)

// WindowFunction is for user interface implements the behaviour
type WindowFunction[IN, OUT any] interface {
	Apply(ctx context.Context, key string, window core.Window, in []IN, out Emitter[OUT]) error
}

type WindowOperator[IN, OUT any] struct {
	assigner    core.WindowAssigner[IN]
	userfunc    WindowFunction[IN, OUT]
	backend     core.StateBackend
	timeSerivce core.InternalTimerService

	windowState core.ValueState[[]IN]
}

func NewWindowOperator[IN, OUT any](assigner core.WindowAssigner[IN], userfunc WindowFunction[IN, OUT]) *WindowOperator[IN, OUT] {
	return &WindowOperator[IN, OUT]{
		assigner:    assigner,
		userfunc:    userfunc,
		timeSerivce: core.NewTimerService(),
	}
}

func (o *WindowOperator[IN, OUT]) Open(ctx context.Context) error {
	be, ok := core.ExtractStateBackend(ctx)
	if !ok {
		return fmt.Errorf("cannot extract state backend, got: %v", be)
	}
	o.backend = be
	windowState, err := be.GetValueState(ctx, "window-elements")
	if err != nil {
		return fmt.Errorf("cannot get value state window-elements error: %w", err)
	}

	o.windowState = core.CastValueType[[]IN](windowState)

	return nil
}

func (o *WindowOperator[IN, OUT]) Process(ctx context.Context, record core.StreamRecord[IN], out Emitter[OUT]) error {
	switch record.Type {
	case core.RecordData:
		// current windows state will hold all record in window
		// assign and get list window
		windows := o.assigner.AssignWindows(record.Value, record.Timestamp)
		for _, w := range windows {
			o.backend.SetCurrentNamespace(w)
			o.backend.SetCurrentKey(record.Key)
			o.timeSerivce.Register(w.MaxTimestamp(), record.Key, w)

			currentWindowState, err := o.windowState.Value(ctx)
			if err != nil {
				return err
			}

			currentWindowState = append(currentWindowState, record.Value)

			o.windowState.Update(ctx, currentWindowState)
		}

		return nil
	case core.RecordWatermark:
		deadlines := o.timeSerivce.AdvanceWatermark(record.Timestamp)

		for _, timer := range deadlines {
			o.backend.SetCurrentKey(timer.Key)
			o.backend.SetCurrentNamespace(timer.Namespace)

			currentWindowState, err := o.windowState.Value(ctx)
			if err != nil {
				return err
			}

			if len(currentWindowState) == 0 {
				continue
			}

			// NOTE: this is workaround create inline window, refactor later

			if err := o.userfunc.Apply(ctx, timer.Key, timer.Namespace.(core.Window), currentWindowState, out); err != nil {
				return err
			}

			if err := o.windowState.Clear(ctx); err != nil {
				return err
			}

		}

		return out.Emit(ctx, core.StreamRecord[OUT]{
			Timestamp: record.Timestamp,
			Type:      core.RecordWatermark,
		})
	default:
		return out.Emit(ctx, core.StreamRecord[OUT]{
			BarrierID: record.BarrierID,
			Timestamp: record.Timestamp,
			Key:       record.Key,
			Type:      record.Type,
		})
	}
}

func (o *WindowOperator[IN, OUT]) Close(ctx context.Context) error {
	return nil
}