package feast

import (
	"context"
	"fmt"
	"time"

	feast "github.com/feast-dev/feast/sdk/go"
	"github.com/feast-dev/feast/sdk/go/protos/feast/serving"
	"github.com/feast-dev/feast/sdk/go/protos/feast/types"
	"go.opentelemetry.io/otel/attribute"
	"go.uber.org/zap"

	"github.com/caraml-dev/merlin/pkg/transformer/spec"
	transTypes "github.com/caraml-dev/merlin/pkg/transformer/types"
	"github.com/caraml-dev/merlin/pkg/transformer/types/converter"
)

type call struct {
	featureTableSpec  *spec.FeatureTable
	columns           []string
	entitySet         map[string]bool
	defaultValues     defaultValues
	columnTypeMapping map[string]types.ValueType_Enum

	feastClient   StorageClient
	servingSource spec.ServingSource

	logger *zap.Logger

	statusMonitoringEnabled bool
	valueMonitoringEnabled  bool
}

func newCall(
	fr *FeastRetriever,
	featureTableSpec *spec.FeatureTable,
	columns []string,
	entitySet map[string]bool,
) (*call, error) {
	feastClient, err := fr.getFeastClient(featureTableSpec.Source)
	if err != nil {
		return nil, err
	}
	return &call{
		featureTableSpec:  featureTableSpec,
		defaultValues:     fr.defaultValues,
		columns:           columns,
		entitySet:         entitySet,
		columnTypeMapping: getFeatureTypeMapping(featureTableSpec),

		feastClient:   feastClient,
		servingSource: featureTableSpec.Source,

		logger: fr.logger,

		statusMonitoringEnabled: fr.options.StatusMonitoringEnabled,
		valueMonitoringEnabled:  fr.options.ValueMonitoringEnabled,
	}, nil
}

// do create request to feast and return the result as table
func (fc *call) do(ctx context.Context, orderedEntityList []orderedFeastRow, features []string) callResult {
	tableName := GetTableName(fc.featureTableSpec)

	feastSource := spec.ServingSource_name[int32(fc.servingSource)]
	ctx, span := tracer.Start(ctx, "feast.doBatchCall")
	span.SetAttributes(attribute.String("feast.source", feastSource))
	span.SetAttributes(attribute.String("table", tableName))
	defer span.End()

	indexes := make([]int, len(orderedEntityList))
	entityList := make([]feast.Row, len(orderedEntityList))
	for i, row := range orderedEntityList {
		entityList[i] = row.Row
		indexes[i] = row.Index
	}
	feastRequest := feast.OnlineFeaturesRequest{
		Project:  fc.featureTableSpec.Project,
		Entities: entityList,
		Features: features,
	}

	startTime := time.Now()
	feastResponse, err := fc.feastClient.GetOnlineFeatures(ctx, &feastRequest)
	durationMs := time.Since(startTime).Milliseconds()
	if err != nil {
		feastLatency.WithLabelValues("error", feastSource).Observe(float64(durationMs))
		feastError.WithLabelValues(feastSource).Inc()

		return callResult{featureTable: nil, err: err}
	}
	feastLatency.WithLabelValues("success", feastSource).Observe(float64(durationMs))
	featureTable, err := fc.processResponse(feastResponse, indexes)
	if err != nil {
		return callResult{featureTable: nil, err: err}
	}

	return callResult{tableName: tableName, featureTable: featureTable, err: nil}
}

func getFeatureTypeMapping(featureTableSpec *spec.FeatureTable) map[string]types.ValueType_Enum {
	mapping := make(map[string]types.ValueType_Enum, len(featureTableSpec.Features))
	for _, feature := range featureTableSpec.Features {
		feastValType := types.ValueType_Enum(types.ValueType_Enum_value[feature.ValueType])
		mapping[feature.Name] = feastValType
	}
	return mapping
}

// processResponse process response from feast serving and create an internal feature table representation of it
func (fc *call) processResponse(feastResponse *feast.OnlineFeaturesResponse, entityIndexes []int) (*internalFeatureTable, error) {
	responseStatus := feastResponse.Statuses()
	responseRows := feastResponse.Rows()
	entities := make([]feast.Row, len(responseRows))
	valueRows := make([]transTypes.ValueRow, len(responseRows))
	indexRows := make([]int, len(responseRows))
	columnTypes := make([]types.ValueType_Enum, len(fc.columns))

	for rowIdx, feastRow := range responseRows {
		valueRow := make(transTypes.ValueRow, len(fc.columns))

		// create entity object, for cache key purpose
		entity := feast.Row{}
		for colIdx, column := range fc.columns {
			var rawValue *types.Value

			featureStatus := responseStatus[rowIdx][column]
			switch featureStatus {
			case serving.FieldStatus_PRESENT:
				rawValue = feastRow[column]
				// set value of entity
				_, isEntity := fc.entitySet[column]
				if isEntity {
					entity[column] = rawValue
				}
			case serving.FieldStatus_NOT_FOUND, serving.FieldStatus_NULL_VALUE, serving.FieldStatus_OUTSIDE_MAX_AGE:
				if columnTypes[colIdx] == types.ValueType_INVALID {
					columnTypes[colIdx] = fc.columnTypeMapping[column]
				}

				defVal, ok := fc.defaultValues.GetDefaultValue(fc.featureTableSpec.Project, column)
				if !ok {
					// no default value is specified, we populate with nil
					valueRow[colIdx] = nil
					continue
				}
				rawValue = defVal

			default:
				return nil, fmt.Errorf("unsupported feature retrieval status for column %s: %s", column, featureStatus)
			}

			val, valType, err := converter.ExtractFeastValue(rawValue)
			if err != nil {
				return nil, err
			}

			// if previously we detected that the column type is invalid then we set it to the correct type
			if valType != types.ValueType_INVALID {
				columnTypes[colIdx] = valType
			}
			valueRow[colIdx] = val

			fc.recordMetrics(val, column, featureStatus)
		}

		entities[rowIdx] = entity
		valueRows[rowIdx] = valueRow
		indexRows[rowIdx] = entityIndexes[rowIdx]
	}

	return &internalFeatureTable{
		entities:    entities,
		columnNames: fc.columns,
		columnTypes: columnTypes,
		valueRows:   valueRows,
		indexRows:   indexRows,
	}, nil
}

func (fc *call) recordMetrics(val interface{}, column string, featureStatus serving.FieldStatus) {
	// put behind feature toggle since it will generate high cardinality metrics
	if fc.valueMonitoringEnabled {
		v, err := converter.ToFloat64(val)
		if err == nil {
			feastFeatureSummary.WithLabelValues(column).Observe(v)
		}
	}

	// put behind feature toggle since it will generate high cardinality metrics
	if fc.statusMonitoringEnabled {
		feastFeatureStatus.WithLabelValues(column, featureStatus.String()).Inc()
	}
}
