package dedup

import (
	"context"
	"sort"

	"github.com/emil/deepres/internal/research/embed"
	"github.com/emil/deepres/internal/research/search"
)

// DedupResult содержит результаты после дедупликации
type DedupResult struct {
	Items       []DedupItem
	RawTexts    []string
}

type DedupItem struct {
	Title       string
	URL         string
	Content     string
	ClusterID   int
}

// Deduplicate принимает результаты поиска, эмбеддит, кластеризует и дедуплицирует
func Deduplicate(ctx context.Context, embedCli *embed.EmbedClient, results []search.SearXNGResult, threshold float32) (*DedupResult, error) {
	if len(results) == 0 {
		return &DedupResult{}, nil
	}

	// Подготавливаем тексты для эмбеддинга
	texts := make([]string, len(results))
	for i, r := range results {
		content := r.Title + ". " + r.Content
		if len([]rune(content)) > 512 {
			content = string([]rune(content)[:512])
		}
		texts[i] = content
	}

	// Эмбеддинги
	embeddings, err := embedCli.Embed(ctx, texts)
	if err != nil {
		return nil, err
	}

	// Простая дедупликация: если косинусная близость > threshold — это дубликат
	n := len(results)
	clusters := make([]int, n)
	clusterID := 0
	for i := range clusters {
		clusters[i] = -1
	}

	for i := 0; i < n; i++ {
		if clusters[i] != -1 {
			continue
		}
		clusters[i] = clusterID
		for j := i + 1; j < n; j++ {
			if clusters[j] != -1 {
				continue
			}
			sim := embed.CosineSimilarity(embeddings[i], embeddings[j])
			if sim >= threshold {
				clusters[j] = clusterID
			}
		}
		clusterID++
	}

	// Собираем результат: по одному лучшему элементу из каждого кластера
	unique := make(map[int]*DedupItem)
	rawTexts := make([]string, 0)

	// Сначала собираем все
	for i, r := range results {
		key := clusters[i]
		if _, exists := unique[key]; !exists {
			unique[key] = &DedupItem{
				Title:   r.Title,
				URL:     r.URL,
				Content: r.Content,
			}
		}
	}

	// Сортируем по clusterID
	keys := make([]int, 0, len(unique))
	for k := range unique {
		keys = append(keys, k)
	}
	sort.Ints(keys)

	var items []DedupItem
	for _, k := range keys {
		item := unique[k]
		items = append(items, *item)
		rawTexts = append(rawTexts, item.Title+". "+item.Content)
	}

	return &DedupResult{
		Items:    items,
		RawTexts: rawTexts,
	}, nil
}