package main

import (
	"context"
	"fmt"
	"strconv"

	"github.com/qdrant/go-client/qdrant"
	"go.uber.org/zap"
)

const collectionName = "flashcard"

type qdrantStore struct {
	client   *qdrant.Client
	embedder *embedder
	sqlite   *sqliteStore
	logger   *zap.Logger
}

func newQdrantStore(client *qdrant.Client, embedder *embedder, sqlite *sqliteStore, logger *zap.Logger) *qdrantStore {
	return &qdrantStore{client: client, embedder: embedder, sqlite: sqlite, logger: logger}
}

func (s *qdrantStore) ensureCollection(ctx context.Context) error {
	exists, err := s.client.CollectionExists(ctx, collectionName)
	if err != nil {
		return fmt.Errorf("check collection exists: %w", err)
	}
	if exists {
		return nil
	}

	err = s.client.CreateCollection(ctx, &qdrant.CreateCollection{
		CollectionName: collectionName,
		VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
			Size:     1536,
			Distance: qdrant.Distance_Cosine,
		}),
	})
	if err != nil {
		return fmt.Errorf("create collection: %w", err)
	}

	return nil
}

func (s *qdrantStore) Index(ctx context.Context, id, publicID int64, version int, note string) error {
	vec, err := s.embedder.embedText(ctx, note)
	if err != nil {
		return fmt.Errorf("embed text: %w", err)
	}

	payload := map[string]*qdrant.Value{
		"_id":       qdrant.NewValueString(strconv.FormatInt(id, 10)),
		"public_id": qdrant.NewValueString(strconv.FormatInt(publicID, 10)),
		"version":   qdrant.NewValueInt(int64(version)),
		"content":   qdrant.NewValueString(note),
	}

	_, err = s.client.Upsert(ctx, &qdrant.UpsertPoints{
		CollectionName: collectionName,
		Points: []*qdrant.PointStruct{
			{
				Id:      qdrant.NewIDNum(uint64(publicID)),
				Vectors: qdrant.NewVectors(vec...),
				Payload: payload,
			},
		},
	})
	if err != nil {
		s.logger.Warn("qdrant index failed (entry in sqlite)", zap.Int64("id", id), zap.Error(err))
	}
	return nil
}

func (s *qdrantStore) Delete(ctx context.Context, id, publicID int64) error {
	wait := true
	_, err := s.client.Delete(ctx, &qdrant.DeletePoints{
		CollectionName: collectionName,
		Wait:           &wait,
		Points: qdrant.NewPointsSelector(
			qdrant.NewIDNum(uint64(publicID)),
		),
	})
	if err != nil {
		s.logger.Warn("qdrant delete failed (tombstone in sqlite)", zap.Int64("public_id", publicID), zap.Error(err))
	}
	return nil
}

func (s *qdrantStore) Search(ctx context.Context, query string, limit uint64) ([]searchResult, error) {
	vec, err := s.embedder.embedText(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("embed query: %w", err)
	}

	if limit == 0 {
		limit = 5
	}

	results, err := s.client.Query(ctx, &qdrant.QueryPoints{
		CollectionName: collectionName,
		Query:          qdrant.NewQueryDense(vec),
		Limit:          &limit,
		WithPayload:    qdrant.NewWithPayload(true),
	})
	if err != nil {
		return nil, fmt.Errorf("query points: %w", err)
	}

	var searchResults []searchResult
	for _, pt := range results {
		r := searchResult{
			Score: pt.Score,
		}

		if payload := pt.GetPayload(); payload != nil {
			if v, ok := payload["_id"]; ok {
				if n, err := strconv.ParseInt(v.GetStringValue(), 10, 64); err == nil {
					r.Entry.ID = n
				}
			}
		}

		if cached, err := s.sqlite.getByID(ctx, r.Entry.ID); err == nil && cached != nil {
			r.Entry = *cached
		}

		searchResults = append(searchResults, r)
	}

	return searchResults, nil
}