package tldrbench

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"
	"unicode"
)

// Fact is one thing a good overview must contain. It passes when a single
// bullet matches every pattern, so a bullet that mentions the date of one
// event next to the name of another does not count. Patterns are
// case-insensitive regular expressions; use "a|b" for alternatives.
type Fact struct {
	Name string   `json:"name"`
	All  []string `json:"all"`
	// Note is free text for you (SuggestFacts puts the chat line it came from
	// here); it is not used for scoring.
	Note string `json:"note,omitempty"`

	res []*regexp.Regexp
}

// LoadFacts reads a JSON list of facts and compiles their patterns.
func LoadFacts(path string) ([]Fact, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var facts []Fact
	if err := json.Unmarshal(raw, &facts); err != nil {
		return nil, fmt.Errorf("%s: %w", path, err)
	}
	for i := range facts {
		if facts[i].Name == "" || len(facts[i].All) == 0 {
			return nil, fmt.Errorf("%s: fact %d needs a name and at least one pattern", path, i+1)
		}
		for _, p := range facts[i].All {
			re, err := regexp.Compile("(?i)" + p)
			if err != nil {
				return nil, fmt.Errorf("%s: fact %q: %w", path, facts[i].Name, err)
			}
			facts[i].res = append(facts[i].res, re)
		}
	}
	return facts, nil
}

// Score is how one overview did. Only the fact counts need a facts file; the
// rest is judged from the overview and the chat alone.
type Score struct {
	FactsHit     int      `json:"factsHit"`
	FactsTotal   int      `json:"factsTotal"`
	Hit          []string `json:"hit,omitempty"`
	Missed       []string `json:"missed,omitempty"`
	Bullets      int      `json:"bullets"`
	LongestChars int      `json:"longestChars"`
	// Invented lists numbers in the overview that appear nowhere in the chat.
	Invented     []string `json:"invented,omitempty"`
	Copied       bool     `json:"copied"`       // looks like the chat pasted back, not a summary
	EnglishIntro bool     `json:"englishIntro"` // opens with "Here is an overview..."
	CutOff       bool     `json:"cutOff"`       // ends mid-sentence, usually the token limit
}

var (
	bulletStart = regexp.MustCompile(`^([-*•]|\d+[.)])\s`)
	numberRe    = regexp.MustCompile(`\d+(?:[.,:]\d+)*`)
	introRe     = regexp.MustCompile(`(?i)^(here( i|')?s|here is|here are|sure|certainly|of course|this is an overview)`)
)

// Bullets groups an overview into its bullets; lines that do not start one
// continue the previous bullet.
func Bullets(out string) []string {
	var bs []string
	for _, ln := range strings.Split(out, "\n") {
		s := strings.TrimSpace(ln)
		if s == "" {
			continue
		}
		if bulletStart.MatchString(s) || len(bs) == 0 {
			bs = append(bs, s)
		} else {
			bs[len(bs)-1] += " " + s
		}
	}
	return bs
}

// normNumber makes "02", "2" and "2,0" compare equal enough for "was this
// number in the chat".
func normNumber(n string) string {
	n = strings.ReplaceAll(n, ",", ".")
	n = strings.TrimLeft(n, "0")
	if n == "" || n[0] == '.' || n[0] == ':' {
		n = "0" + n
	}
	return n
}

func numbers(s string) map[string]bool {
	set := map[string]bool{}
	for _, n := range numberRe.FindAllString(s, -1) {
		set[normNumber(n)] = true
		// "14:30" also makes "14" and "30" legitimate on their own.
		for _, part := range regexp.MustCompile(`[.,:]`).Split(n, -1) {
			set[normNumber(part)] = true
		}
	}
	return set
}

// Evaluate scores an overview. chat is the transcript the model was given.
func Evaluate(out, chat string, facts []Fact) Score {
	bullets := Bullets(out)
	var sc Score
	sc.Bullets = len(bullets)
	for _, b := range bullets {
		if n := len([]rune(b)); n > sc.LongestChars {
			sc.LongestChars = n
		}
	}

	sc.FactsTotal = len(facts)
	for _, f := range facts {
		ok := false
		for _, b := range bullets {
			all := true
			for _, re := range f.res {
				if !re.MatchString(b) {
					all = false
					break
				}
			}
			if all {
				ok = true
				break
			}
		}
		if ok {
			sc.FactsHit++
			sc.Hit = append(sc.Hit, f.Name)
		} else {
			sc.Missed = append(sc.Missed, f.Name)
		}
	}

	known := numbers(chat)
	seen := map[string]bool{}
	for _, b := range bullets {
		b = bulletStart.ReplaceAllString(b, "")
		for _, n := range numberRe.FindAllString(b, -1) {
			if !known[normNumber(n)] && !seen[n] {
				seen[n] = true
				sc.Invented = append(sc.Invented, n)
			}
		}
	}

	trimmed := strings.TrimSpace(out)
	sc.Copied = (sc.Bullets <= 2 && len([]rune(trimmed)) > 600) || sc.LongestChars > 700
	sc.EnglishIntro = introRe.MatchString(trimmed)
	if r := []rune(trimmed); len(r) > 0 {
		last := r[len(r)-1]
		sc.CutOff = unicode.IsLetter(last) || unicode.IsDigit(last) || last == ',' || last == ';' || last == ':'
	}
	return sc
}