package kjv

import (
	"sort"
	"strings"
	"unicode"
)

// Speller checks and suggests spellings from kjv.dic (a plain word list) and
// kjv.aff (which only adds a TRY order and REP replacement hints), so no C
// library is needed.
type Speller struct {
	words   map[string]bool     // exact entries
	byLower map[string][]string // lowercase -> entries
	list    []string            // lowercase entries, dictionary order, unique
	try     []rune
	rep     [][2]string
}

func NewSpeller(dic, aff []byte) *Speller {
	s := &Speller{words: map[string]bool{}, byLower: map[string][]string{}}
	lines := strings.Split(string(dic), "\n")
	seen := map[string]bool{}
	for i, line := range lines {
		line = strings.TrimSpace(line)
		if i == 0 || line == "" { // first line is the entry count
			continue
		}
		if k := strings.IndexByte(line, '/'); k >= 0 {
			line = line[:k]
		}
		s.words[line] = true
		low := strings.ToLower(line)
		s.byLower[low] = append(s.byLower[low], line)
		if !seen[low] {
			seen[low] = true
			s.list = append(s.list, low)
		}
	}
	for _, line := range strings.Split(string(aff), "\n") {
		f := strings.Fields(line)
		switch {
		case len(f) == 2 && f[0] == "TRY":
			s.try = []rune(f[1])
		case len(f) == 3 && f[0] == "REP":
			s.rep = append(s.rep, [2]string{f[1], strings.ReplaceAll(f[2], "_", " ")})
		}
	}
	return s
}

func stripSyntax(w string) string {
	return strings.Map(func(r rune) rune {
		if r == '+' || r == '^' || r == '*' {
			return -1
		}
		return r
	}, w)
}

func hasLetter(w string) bool {
	for _, r := range w {
		if unicode.IsLetter(r) {
			return true
		}
	}
	return false
}

func (s *Speller) Spell(word string) bool {
	word = stripSyntax(word)
	if !hasLetter(word) || s.words[word] {
		return true
	}
	low := strings.ToLower(word)
	if low != word && (word == strings.ToUpper(word) || word == capitalize(low)) {
		return len(s.byLower[low]) > 0
	}
	return false
}

func capitalize(w string) string {
	r := []rune(w)
	if len(r) == 0 {
		return w
	}
	return strings.ToUpper(string(r[:1])) + strings.ToLower(string(r[1:]))
}

func (s *Speller) known(low string) []string { return s.byLower[low] }

// Suggest returns likely corrections, best first.
func (s *Speller) Suggest(word string) []string {
	word = stripSyntax(word)
	low := strings.ToLower(word)
	if low == "" {
		return nil
	}
	var out []string
	seen := map[string]bool{}
	add := func(cands ...string) {
		for _, c := range cands {
			if !seen[c] {
				seen[c] = true
				out = append(out, c)
			}
		}
	}
	addLower := func(l string) { // dictionary form(s) of a lowercase candidate
		if l == low {
			return
		}
		add(s.known(l)...)
	}
	if low != word { // wrong capitalization of a known word
		add(s.known(low)...)
	}
	// replacement hints from the .aff file, applied at each occurrence
	for _, r := range s.rep {
		for at := 0; ; {
			i := strings.Index(low[at:], r[0])
			if i < 0 {
				break
			}
			i += at
			cand := low[:i] + r[1] + low[i+len(r[0]):]
			if strings.Contains(cand, " ") {
				if p := strings.Fields(cand); len(p) == 2 && len(s.known(p[0])) > 0 && len(s.known(p[1])) > 0 {
					add(p[0] + " " + p[1])
				}
			} else {
				addLower(cand)
			}
			at = i + 1
		}
	}
	rs := []rune(low)
	// one edit away: swap, delete, replace, insert
	for i := 0; i+1 < len(rs); i++ {
		c := append([]rune(nil), rs...)
		c[i], c[i+1] = c[i+1], c[i]
		addLower(string(c))
	}
	for i := range rs {
		addLower(string(rs[:i]) + string(rs[i+1:]))
	}
	for i := range rs {
		for _, t := range s.try {
			addLower(string(rs[:i]) + string(t) + string(rs[i+1:]))
		}
	}
	for i := 0; i <= len(rs); i++ {
		for _, t := range s.try {
			addLower(string(rs[:i]) + string(t) + string(rs[i:]))
		}
	}
	// two words run together ("lordgod")
	for i := 1; i < len(rs); i++ {
		a, b := string(rs[:i]), string(rs[i:])
		if len(a) > 1 && len(b) > 1 && len(s.known(a)) > 0 && len(s.known(b)) > 0 {
			add(a + " " + b)
		}
	}
	// still little to offer: anything within two edits, closest first
	if len(out) < 3 {
		type cand struct {
			w    string
			dist int
			pre  int
		}
		var cs []cand
		for _, w := range s.list {
			if abs(len(w)-len(low)) > 2 || w == low {
				continue
			}
			if d := editDistance(low, w, 2); d <= 2 {
				cs = append(cs, cand{w, d, commonPrefix(low, w)})
			}
		}
		sort.SliceStable(cs, func(i, j int) bool {
			if cs[i].dist != cs[j].dist {
				return cs[i].dist < cs[j].dist
			}
			return cs[i].pre > cs[j].pre
		})
		for _, c := range cs {
			addLower(c.w)
			if len(out) >= 8 {
				break
			}
		}
	}
	if len(out) > 10 {
		out = out[:10]
	}
	if word == strings.ToUpper(word) && len(word) > 1 {
		for i := range out {
			out[i] = strings.ToUpper(out[i])
		}
	} else if word == capitalize(low) && word != low {
		for i := range out {
			out[i] = capitalize(out[i])
		}
	}
	return out
}

func abs(n int) int {
	if n < 0 {
		return -n
	}
	return n
}

func commonPrefix(a, b string) int {
	n := 0
	for n < len(a) && n < len(b) && a[n] == b[n] {
		n++
	}
	return n
}

// editDistance is optimal-string-alignment distance, giving up above limit.
func editDistance(a, b string, limit int) int {
	ra, rb := []rune(a), []rune(b)
	prev2 := make([]int, len(rb)+1)
	prev := make([]int, len(rb)+1)
	cur := make([]int, len(rb)+1)
	for j := range prev {
		prev[j] = j
	}
	for i := 1; i <= len(ra); i++ {
		cur[0] = i
		best := cur[0]
		for j := 1; j <= len(rb); j++ {
			cost := 1
			if ra[i-1] == rb[j-1] {
				cost = 0
			}
			v := min(prev[j]+1, cur[j-1]+1, prev[j-1]+cost)
			if i > 1 && j > 1 && ra[i-1] == rb[j-2] && ra[i-2] == rb[j-1] {
				v = min(v, prev2[j-2]+1)
			}
			cur[j] = v
			best = min(best, v)
		}
		if best > limit {
			return limit + 1
		}
		prev2, prev, cur = prev, cur, prev2
	}
	return prev[len(rb)]
}