package search

import (
	"searchEngine/analyzer"
	"searchEngine/index"
	"slices"
)

type Result struct {
	DocID int
	Score int
}

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()

	firstToken := tokenizedQuery[0]
	basePostings, exists := idx.Data[firstToken]
	if !exists {
		return nil
	}

	scores := make(map[int]int)
	for _, p := range basePostings {
		scores[p.DocId] = p.Count
	}

	for i := 1; i < len(tokenizedQuery); i++ {
		token := tokenizedQuery[i]
		nextPostings, exists := idx.Data[token]
		if !exists {
			return nil
		}

		lookup := make(map[int]int)
		for _, p := range nextPostings {
			lookup[p.DocId] = p.Count
		}

		newScores := make(map[int]int)
		for docID, currentScore := range scores {
			if nextCount, found := lookup[docID]; found {
				newScores[docID] = currentScore + nextCount
			}
		}
		scores = newScores
	}

	var results []Result
	for docID, score := range scores {
		results = append(results, Result{DocID: docID, Score: score})
	}

	slices.SortFunc(results, func(a, b Result) int {
		return 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
}