package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"

	"github.com/google/uuid"
	"github.com/qdrant/go-client/qdrant"
)

const collectionName = "knowledge"

type qdrantStore struct {
	client   *qdrant.Client
	embedder *embedder
}

func newQdrantStore(client *qdrant.Client, embedder *embedder) *qdrantStore {
	return &qdrantStore{client: client, embedder: embedder}
}

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, // text-embedding-3-small dimension
			Distance: qdrant.Distance_Cosine,
		}),
	})
	if err != nil {
		return fmt.Errorf("create collection: %w", err)
	}

	return nil
}

type storedKnowledge struct {
	ID       string         `json:"id"`
	Content  string         `json:"content"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

type searchResult struct {
	ID       string         `json:"id"`
	Content  string         `json:"content"`
	Score    float32        `json:"score"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

func (s *qdrantStore) store(ctx context.Context, content string, metadata map[string]any) (*storedKnowledge, error) {
	vec, err := s.embedder.embedText(ctx, content)
	if err != nil {
		return nil, fmt.Errorf("embed text: %w", err)
	}

	id := uuid.New().String()
	payload := map[string]*qdrant.Value{
		"content": qdrant.NewValueString(content),
	}

	if metadata != nil {
		metaBytes, err := json.Marshal(metadata)
		if err != nil {
			return nil, fmt.Errorf("marshal metadata: %w", err)
		}
		payload["_metadata"] = qdrant.NewValueString(string(metaBytes))
	}

	idHash := sha256.Sum256([]byte(id))
	idNum := uint64(0)
	for i := 0; i < 8; i++ {
		idNum = (idNum << 8) | uint64(idHash[i])
	}

	_, err = s.client.Upsert(ctx, &qdrant.UpsertPoints{
		CollectionName: collectionName,
		Points: []*qdrant.PointStruct{
			{
				Id:      qdrant.NewIDNum(idNum),
				Vectors: qdrant.NewVectors(vec...),
				Payload: payload,
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("upsert point: %w", err)
	}

	return &storedKnowledge{
		ID:       id,
		Content:  content,
		Metadata: metadata,
	}, 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["content"]; ok {
				r.ID = fmt.Sprintf("%d", pt.GetId().GetNum())
				r.Content = v.GetStringValue()
			}
			if v, ok := payload["_metadata"]; ok {
				var meta map[string]any
				if err := json.Unmarshal([]byte(v.GetStringValue()), &meta); err == nil {
					r.Metadata = meta
				}
			}
		}

		searchResults = append(searchResults, r)
	}

	return searchResults, nil
}