search.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/blevesearch/bleve"
	"graph"
	"net/http"
	"strings"
)

func SkipName(name string) bool {
	return strings.HasSuffix(name, ".jpg") ||
		strings.HasSuffix(name, ".jpeg") ||
		strings.HasSuffix(name, ".png") ||
		strings.HasSuffix(name, ".gif") ||
		strings.HasSuffix(name, ".gifv")
}

func Rebuild(index bleve.Index) {
	var x map[string]interface{}
	graph.CayleyGremlin(&x, "g.V().All()")

	for _, node := range x["result"].([]interface{}) {
		id := node.(map[string]interface{})["id"].(string)
		if SkipName(id) {
			continue
		}
		fmt.Printf("\radding '\x1b[34;1m%v\x1b[0m' ... ", id)
		index.Index(id, node.(map[string]interface{}))
		fmt.Println("OK")
	}
}

func Search(input string) (*bleve.SearchResult, error) {
	query := bleve.NewBooleanQuery(nil, nil, nil)
	for _, term := range strings.Split(input, " ") {
		query.AddShould(bleve.NewFuzzyQuery(term))
		query.AddShould(bleve.NewTermQuery(term))
	}
	searchRequest := bleve.NewSearchRequest(query)
	return index.Search(searchRequest)
}

var index bleve.Index

type Result struct {
	ID    string  `json:"id"`
	Score float64 `json:"score"`
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	term := r.FormValue("q")
	searchResult, _ := Search(term)
	results := []Result{}
	for _, doc := range searchResult.Hits {
		if doc.Score > 0.1 {
			results = append(results, Result{ID: doc.ID, Score: doc.Score})
		}
	}
	res := map[string]interface{}{
		"query":  term,
		"result": results,
	}
	res_json, _ := json.Marshal(res)
	fmt.Fprintf(w, string(res_json))
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "UPDATE" {
		Rebuild(index)
		rebuild_res := map[string]interface{}{
			"err": nil,
			"msg": "rebuilt index",
		}
		rebuild_data, _ := json.Marshal(rebuild_res)
		fmt.Fprintf(w, string(rebuild_data))
	}
}

func main() {
	var rebuild_index bool
	var web_daemon bool
	flag.BoolVar(&rebuild_index, "rebuild", false, "rebuild index")
	flag.BoolVar(&web_daemon, "web", false, "start http server")
	flag.Parse()

	mapping := bleve.NewIndexMapping()
	var err error
	index, err = bleve.New("hb.bleve", mapping)
	if err == bleve.ErrorIndexPathExists {
		index, _ = bleve.Open("hb.bleve")
	}

	if rebuild_index {
		fmt.Println("rebuilding index")
		Rebuild(index)
	}

	if web_daemon {
		fmt.Println("starting web daemon")
		http.Handle("/query", http.HandlerFunc(searchHandler))
		http.Handle("/rebuild", http.HandlerFunc(indexHandler))
		err := http.ListenAndServe(":8282", nil)
		fmt.Println("ended with:", err)
	} else {
		term := strings.Join(flag.Args(), " ")
		searchResult, _ := Search(term)
		fmt.Println(searchResult)
	}
}