package core

import (
    "github.com/huichen/wukong/types"
    "github.com/huichen/wukong/utils"
    "log"
    "sort"
    "sync"
)

type Ranker struct {
    lock struct {
        sync.RWMutex
        fields map[uint64]interface{}
        docs   map[uint64]bool
    }
    initialized bool
}

func (ranker *Ranker) Init() {
    if ranker.initialized == true {
        log.Fatal("排序器不能初始化两次")
    }
    ranker.initialized = true

    ranker.lock.fields = make(map[uint64]interface{})
    ranker.lock.docs = make(map[uint64]bool)
}

// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(docId uint64, fields interface{}) {
    if ranker.initialized == false {
        log.Fatal("排序器尚未初始化")
    }

    ranker.lock.Lock()
    ranker.lock.fields[docId] = fields
    ranker.lock.docs[docId] = true
    ranker.lock.Unlock()
}

// 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId uint64) {
    if ranker.initialized == false {
        log.Fatal("排序器尚未初始化")
    }

    ranker.lock.Lock()
    delete(ranker.lock.fields, docId)
    delete(ranker.lock.docs, docId)
    ranker.lock.Unlock()
}

// 给文档评分并排序
func (ranker *Ranker) Rank(
    docs []types.IndexedDocument, options types.RankOptions, countDocsOnly bool) (types.ScoredDocuments, int) {
    if ranker.initialized == false {
        log.Fatal("排序器尚未初始化")
    }

    // 对每个文档评分
    var outputDocs types.ScoredDocuments
    numDocs := 0
    for _, d := range docs {
        ranker.lock.RLock()
        // 判断doc是否存在
        if _, ok := ranker.lock.docs[d.DocId]; ok {
            fs := ranker.lock.fields[d.DocId]
            ranker.lock.RUnlock()
            // 计算评分并剔除没有分值的文档
            scores := options.ScoringCriteria.Score(d, fs)
            if len(scores) > 0 {
                if !countDocsOnly {
                    outputDocs = append(outputDocs, types.ScoredDocument{
                        DocId:                 d.DocId,
                        Scores:                scores,
                        TokenSnippetLocations: d.TokenSnippetLocations,
                        TokenLocations:        d.TokenLocations})
                }
                numDocs++
            }
        } else {
            ranker.lock.RUnlock()
        }
    }

    // 排序
    if !countDocsOnly {
        if options.ReverseOrder {
            sort.Sort(sort.Reverse(outputDocs))
        } else {
            sort.Sort(outputDocs)
        }
        // 当用户要求只返回部分结果时返回部分结果
        var start, end int
        if options.MaxOutputs != 0 {
            start = utils.MinInt(options.OutputOffset, len(outputDocs))
            end = utils.MinInt(options.OutputOffset+options.MaxOutputs, len(outputDocs))
        } else {
            start = utils.MinInt(options.OutputOffset, len(outputDocs))
            end = len(outputDocs)
        }
        return outputDocs[start:end], numDocs
    }
    return outputDocs, numDocs
}