package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"text/template"
)

var (
	token     = flag.String("token", "", "API-Token to use instead of $FASTGPT_API_TOKEN or ~/fastgpt_token")
	stdin     = flag.Bool("stdin", true, "Use stdin instead of command-line arguments for search")
	cache     = flag.Bool("cache", false, "Whether to allow cached requests & responses.")
	websearch = flag.Bool("websearch", true, "Whether to perform web searches to enrich answers. MUST   be true for now.")
	raw       = flag.Bool("raw", false, "Output raw json result")
)

type Query struct {
	Query     string `json:"query"`      // A query to be answered.
	Cache     bool   `json:"cache"`      // Whether to allow cached requests & responses. (default true)
	WebSearch bool   `json:"web_search"` // Whether to perform web searches to enrich answers. MUST be true for now.
}

type Result struct {
	Meta struct {
		Id   string `json:"id"`
		Node string `json:"node"`
		MS   int    `json:"ms"`
	} `json:"meta"`
	Data struct {
		Query      string      `json:"-"`
		Output     string      `json:"output"`     // Answer output
		References []Reference `json:"references"` // The search results that are referred in the answer.
		Tokens     int         `json:"tokens"`     // Amount of tokens processed
	} `json:"data"`
}

type Reference struct {
	Title   string `json:"title"`   // Title of the referenced search result.
	Snippet string `json:"snippet"` // Snippet of the referenced search result.
	URL     string `json:"url"`     // URL of the referenced search result.
}

var fm = template.FuncMap{
	"quote": func(s string) string {
		buf := &bytes.Buffer{}
		scn := bufio.NewScanner(strings.NewReader(s))
		for scn.Scan() {
			buf.WriteString("> ")
			buf.WriteString(scn.Text())
			buf.WriteString("\n")
		}
		return buf.String()
	},
	"plusone": func(i int) int { return i + 1 },
}
var tmpl = template.Must(template.New("output").
	Funcs(fm).
	Parse(`Q: {{ .Query }}

{{.Output}}
{{ with .References }}{{- range $i, $ref := . }}
[{{ plusone $i }}]: [{{$ref.Title}}]({{$ref.URL}})
{{ quote $ref.Snippet }}
{{- end }}
{{- end }}
`))

func main() {

	q := Query{
		Cache:     *cache,
		WebSearch: *websearch,
	}
	buf := &bytes.Buffer{}

	flag.Parse()

	if *token == "" {
		*token = os.Getenv("FASTGPT_API_TOKEN")
		if *token == "" {
			h, _ := os.UserHomeDir()
			b, e := os.ReadFile(h + "/.fastgpt_token")
			if e != nil {
				panic(e)
			}
			*token = string(b)
		}
	}

	if *token == "" {
		panic("token needed")
	}

	if *stdin {
		if _, e := io.Copy(buf, os.Stdin); e != nil {
			panic(e)
		}
		q.Query = buf.String()
		buf.Reset()
	} else {
		q.Query = strings.Join(os.Args, " ")
	}

	if e := json.NewEncoder(buf).Encode(&q); e != nil {
		panic(e)
	}

	req, e := http.NewRequest("POST",
		"https://kagi.com/api/v0/fastgpt",
		buf)
	if e != nil {
		panic(e)
	}
	req.Header.Add("Authorization", "Bot "+*token)
	req.Header.Add("Content-Type", "application/json")

	r, e := http.DefaultClient.Do(req)
	if e != nil {
		panic(e)
	}

	if r.StatusCode != 200 {
		log.Printf("Non-200 status code: %d\nBody: \n", r.StatusCode)
		io.Copy(os.Stdout, r.Body)
		os.Exit(r.StatusCode)
	}

	buf.Reset()
	res := &Result{}
	if e := json.NewDecoder(io.TeeReader(r.Body, buf)).Decode(res); e != nil {
		log.Printf("Couldn't decode result: %v\nBody:\n", e)
		io.Copy(os.Stdout, buf)
	}
	res.Data.Query = q.Query

	if *raw {
		io.Copy(os.Stdout, buf)
	} else {
		tmpl.Execute(os.Stdout, res.Data)
	}
}