package vector

import (
	"context"
	"crypto/tls"
	"fmt"
	"moss-dev/gpt"

	"github.com/google/uuid"
	"github.com/juju/errors"
	pb "github.com/qdrant/go-client/qdrant"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
)

type QdrantVectorEngine struct {
	endpoint string
	token    string
	ai       *gpt.GPT

	conn *grpc.ClientConn
}

func CreateQdrantVectorEngine(endpoint, token string, ai *gpt.GPT) (*QdrantVectorEngine, error) {
	// Set up a connection to the server.
	config := &tls.Config{}
	conn, err := grpc.Dial(endpoint, grpc.WithTransportCredentials(credentials.NewTLS(config)))
	if err != nil {
		return nil, errors.Errorf("create qdrant vector engine grpc connect, endpoint: %s, err: %v", endpoint, err)
	}

	return &QdrantVectorEngine{endpoint: endpoint, token: token, ai: ai, conn: conn}, nil
}

func (q *QdrantVectorEngine) UpsertTableSchema(ctx context.Context, collectName string, project, db, table, schema string) (string, error) {
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	embeddings, err := q.ai.CreateEmbedding(ctx, []string{schema})
	if err != nil {
		return "", errors.Errorf("fail to create embedding %v", err)
	}

	pointID := q.pointID(project, db, table, "schema")
	wait := true
	pointCli := pb.NewPointsClient(q.conn)
	resp, err := pointCli.Upsert(ctx, &pb.UpsertPoints{
		CollectionName: collectName,
		Wait:           &wait,
		Points: []*pb.PointStruct{
			&pb.PointStruct{
				Id: &pb.PointId{PointIdOptions: &pb.PointId_Uuid{Uuid: pointID}},
				Payload: map[string]*pb.Value{
					"project": &pb.Value{Kind: &pb.Value_StringValue{StringValue: project}},
					"db":      &pb.Value{Kind: &pb.Value_StringValue{StringValue: db}},
					"table":   &pb.Value{Kind: &pb.Value_StringValue{StringValue: table}},
					"schema":  &pb.Value{Kind: &pb.Value_StringValue{StringValue: schema}},
				},
				Vectors: &pb.Vectors{
					VectorsOptions: &pb.Vectors_Vector{Vector: &pb.Vector{Data: embeddings[0]}},
				},
			},
		},
	})
	if err != nil {
		return "", errors.Trace(err)
	}

	if resp.Result.GetStatus() == pb.UpdateStatus_UnknownUpdateStatus {
		return pointID, errors.Errorf("upsert schema into verctor engine: unknown update status")
	}

	return pointID, nil
}

func (q *QdrantVectorEngine) GetCollection(ctx context.Context, name string) (string, error) {
	// Contact the server and print out its response.
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	collectionCli := pb.NewCollectionsClient(q.conn)
	resp, err := collectionCli.Get(ctx, &pb.GetCollectionInfoRequest{
		CollectionName: name,
	})
	if err != nil {
		return "", errors.Trace(err)
	}

	return resp.String(), nil
}

func (q *QdrantVectorEngine) CreateCollection(ctx context.Context, name string) error {
	// Contact the server and print out its response.
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	collectionCli := pb.NewCollectionsClient(q.conn)
	resp, err := collectionCli.Create(ctx, &pb.CreateCollection{
		CollectionName: name,
		VectorsConfig: &pb.VectorsConfig{Config: &pb.VectorsConfig_Params{
			Params: &pb.VectorParams{
				Size:     1536, // only support openai.text-embedding-ada-002 temporially
				Distance: pb.Distance_Cosine,
			},
		}},
	})
	if err != nil {
		return errors.Trace(err)
	}

	if !resp.Result {
		return errors.Errorf("fail to create collection %s", name)
	}

	return nil
}

func (q *QdrantVectorEngine) DeleteCollection(ctx context.Context, name string) error {
	// Contact the server and print out its response.
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	collectionCli := pb.NewCollectionsClient(q.conn)
	resp, err := collectionCli.Delete(ctx, &pb.DeleteCollection{
		CollectionName: name,
	})
	if err != nil {
		return errors.Trace(err)
	}

	if !resp.Result {
		return errors.Errorf("fail to delete collection %s", name)
	}

	return nil
}

func (q *QdrantVectorEngine) GetVector(ctx context.Context, collectName, project, schema, table string) (string, string, []float32, error) {
	// Contact the server and print out its response.
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	pointID := q.pointID(project, schema, table, "schema")
	pointCli := pb.NewPointsClient(q.conn)
	resp, err := pointCli.Get(ctx, &pb.GetPoints{
		CollectionName: collectName,
		Ids:            []*pb.PointId{&pb.PointId{PointIdOptions: &pb.PointId_Uuid{Uuid: pointID}}},
		WithPayload:    &pb.WithPayloadSelector{SelectorOptions: &pb.WithPayloadSelector_Enable{Enable: true}},
		WithVectors:    &pb.WithVectorsSelector{SelectorOptions: &pb.WithVectorsSelector_Enable{Enable: true}},
	})
	if err != nil {
		return "", "", nil, errors.Trace(err)
	}

	if len(resp.Result) < 1 {
		return "", "", nil, errors.NotFoundf("vector of %s:%s:%s", project, schema, table)
	}

	res := resp.Result[0]
	return res.Id.GetUuid(), res.Payload["schema"].GetStringValue(), res.Vectors.GetVector().GetData(), nil
}

func (q *QdrantVectorEngine) Search(ctx context.Context, collectName, query string) ([]string, []string, [][]float32, []float32, error) {
	// Contact the server and print out its response.
	md := metadata.New(map[string]string{"api-key": q.token})
	ctx = metadata.NewOutgoingContext(ctx, md)

	embeddings, err := q.ai.CreateEmbedding(ctx, []string{query})
	if err != nil {
		return nil, nil, nil, nil, errors.Errorf("fail to create embedding %v", err)
	}

	pointCli := pb.NewPointsClient(q.conn)
	resp, err := pointCli.Search(ctx, &pb.SearchPoints{
		CollectionName: collectName,
		Vector:         embeddings[0],
		Limit:          1,
		WithPayload:    &pb.WithPayloadSelector{SelectorOptions: &pb.WithPayloadSelector_Enable{Enable: true}},
		WithVectors:    &pb.WithVectorsSelector{SelectorOptions: &pb.WithVectorsSelector_Enable{Enable: true}},
	})
	if err != nil {
		return nil, nil, nil, nil, errors.Trace(err)
	}

	resPointID := make([]string, 0, 3)
	resSchema := make([]string, 0, 3)
	resVector := make([][]float32, 0, 3)
	resScore := make([]float32, 0, 3)

	for _, res := range resp.Result {
		resPointID = append(resPointID, res.Id.GetUuid())
		resSchema = append(resSchema, res.Payload["schema"].GetStringValue())
		resVector = append(resVector, res.Vectors.GetVector().GetData())
		resScore = append(resScore, res.Score)
	}
	return resPointID, resSchema, resVector, resScore, nil
}

func (q *QdrantVectorEngine) Close() error {
	return q.conn.Close()
}

func (q *QdrantVectorEngine) pointID(project, db, table, category string) string {
	name := fmt.Sprintf("%s:%s:%s:%s", project, db, table, category)
	v5 := uuid.NewSHA1(uuid.NameSpaceURL, []byte(name))
	return v5.String()
}