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)
	docPostings := make(map[int][]index.Posting)
	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

			docPostings[p.DocId] = append(docPostings[p.DocId], p)
		}
	}

	var results []Result
	for docID, count := range docMatchCounts {
		if count == len(uniqueTokens) {
			results = append(results, Result{DocID: docID, Score: scores[docID]})
		}
	}
	for i := range results {
		proximityBonus := 0.0
		resultsPostings := docPostings[results[i].DocID]
		for j := 0; j < len(resultsPostings)-1; j++ {
			bodyBonus := scoring.CalculateProximity(resultsPostings[j].BodyPositions, resultsPostings[j+1].BodyPositions)

			titleBonus := scoring.CalculateProximity(resultsPostings[j].TitlePositions, resultsPostings[j+1].TitlePositions)

			proximityBonus += bodyBonus + (titleBonus * 5.0)
		}

		results[i].Score += proximityBonus
	}
	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
}