package search

import (
	"cmp"
	"searchEngine/core/analyzer"
	"searchEngine/core/index"
	"searchEngine/core/scoring"
	"slices"
)

type Result struct {
	DocID int
	Score float64
}

func Search(idx *index.InvertedIndex, query string) []string {
	tokenizedQuery := analyzer.ProcessText(query)
	if len(tokenizedQuery) == 0 {
		return nil
	}

	idx.Mu.RLock()
	defer idx.Mu.RUnlock()

	totalDocs := len(idx.DocNames)
	docMatchCounts := make(map[int]int)
	scores := make(map[int]float64)

	uniqueMap := make(map[string]bool)
	var uniqueTokens []string
	for _, token := range tokenizedQuery {
		if !uniqueMap[token] {
			uniqueMap[token] = true
			uniqueTokens = append(uniqueTokens, token)
		}
	}

	for _, token := range tokenizedQuery {
		postings, exists := idx.Data[token]
		if !exists {
			return nil
		}

		docFrequency := len(postings)
		idf := scoring.CalculateIDF(totalDocs, docFrequency)

		for _, p := range postings {
			docMatchCounts[p.DocId]++
			termScore := scoring.ScoreTFIDF(p, idf)
			scores[p.DocId] += termScore
		}
	}

	var results []Result
	for docID, count := range docMatchCounts {
		if count == len(uniqueTokens) {
			results = append(results, Result{DocID: docID, Score: scores[docID]})
		}
	}

	slices.SortFunc(results, func(a, b Result) int {
		return cmp.Compare(b.Score, a.Score)
	})

	var titles []string
	for _, res := range results {
		if title, exists := idx.DocNames[res.DocID]; exists {
			titles = append(titles, title)
		}
	}

	return titles
}