package store

import (
	"context"
	"fmt"

	grpc "github.com/mudler/LocalAI/pkg/grpc"
	"github.com/mudler/LocalAI/pkg/grpc/proto"
)

// Wrapper for the GRPC client so that simple use cases are handled without verbosity

// SetCols sets multiple key-value pairs in the store
// It's in columnar format so that keys[i] is associated with values[i]
func SetCols(ctx context.Context, c grpc.Backend, keys [][]float32, values [][]byte) error {
	res, err := c.StoresSet(ctx, &proto.StoresSetOptions{
		Keys:   WrapKeys(keys),
		Values: WrapValues(values),
	})
	if err != nil {
		return err
	}

	if res.Success {
		return nil
	}

	return fmt.Errorf("failed to set keys: %v", res.Message)
}

// SetSingle sets a single key-value pair in the store
// Don't call this in a tight loop, instead use SetCols
func SetSingle(ctx context.Context, c grpc.Backend, key []float32, value []byte) error {
	return SetCols(ctx, c, [][]float32{key}, [][]byte{value})
}

// DeleteCols deletes multiple key-value pairs from the store
// It's in columnar format so that keys[i] is associated with values[i]
func DeleteCols(ctx context.Context, c grpc.Backend, keys [][]float32) error {
	res, err := c.StoresDelete(ctx, &proto.StoresDeleteOptions{Keys: WrapKeys(keys)})
	if err != nil {
		return err
	}

	if res.Success {
		return nil
	}

	return fmt.Errorf("failed to delete keys: %v", res.Message)
}

// DeleteSingle deletes a single key-value pair from the store
// Don't call this in a tight loop, instead use DeleteCols
func DeleteSingle(ctx context.Context, c grpc.Backend, key []float32) error {
	return DeleteCols(ctx, c, [][]float32{key})
}

// GetCols gets multiple key-value pairs from the store
// It's in columnar format so that keys[i] is associated with values[i]
// Be warned the keys are sorted and will be returned in a different order than they were input
// There is no guarantee as to how the keys are sorted
func GetCols(ctx context.Context, c grpc.Backend, keys [][]float32) ([][]float32, [][]byte, error) {
	res, err := c.StoresGet(ctx, &proto.StoresGetOptions{Keys: WrapKeys(keys)})
	if err != nil {
		return nil, nil, err
	}
	return UnwrapKeys(res.Keys), UnwrapValues(res.Values), nil
}

// GetSingle gets a single key-value pair from the store
// Don't call this in a tight loop, instead use GetCols
func GetSingle(ctx context.Context, c grpc.Backend, key []float32) ([]byte, error) {
	_, values, err := GetCols(ctx, c, [][]float32{key})
	if err != nil {
		return nil, err
	}

	if len(values) > 0 {
		return values[0], nil
	}

	return nil, fmt.Errorf("failed to get key")
}

// Find similar keys to the given key. Returns the keys, values, and similarities
func Find(ctx context.Context, c grpc.Backend, key []float32, topk int) ([][]float32, [][]byte, []float32, error) {
	res, err := c.StoresFind(ctx, &proto.StoresFindOptions{
		Key:  &proto.StoresKey{Floats: key},
		TopK: int32(topk),
	})
	if err != nil {
		return nil, nil, nil, err
	}
	return UnwrapKeys(res.Keys), UnwrapValues(res.Values), res.Similarities, nil
}