package wordsearchii

type State struct {
	X, Y     int
	Previous *State
}

func (s *State) Visited(x, y int) bool {
	if s == nil {
		return false
	}

	if s.X == x && s.Y == y {
		return true
	}

	return s.Previous.Visited(x, y)
}

type Trie struct {
	Children map[byte]*Trie
	States   []*State
}

func NewTrie(board [][]byte) *Trie {
	trie := &Trie{
		Children: map[byte]*Trie{},
		States:   []*State{},
	}

	for x := range board {
		for y := range board[x] {
			if _, ok := trie.Children[board[x][y]]; !ok {
				trie.Children[board[x][y]] = &Trie{
					Children: map[byte]*Trie{},
					States:   []*State{},
				}
			}

			trie.Children[board[x][y]].States = append(trie.Children[board[x][y]].States, &State{
				X:        x,
				Y:        y,
				Previous: nil,
			})
		}
	}

	return trie
}

var CardinalDirections = [][2]int{
	{-1, 0},
	{1, 0},
	{0, -1},
	{0, 1},
}

func (t *Trie) Explore(board [][]byte, ch byte) *Trie {
	child := &Trie{
		Children: map[byte]*Trie{},
		States:   []*State{},
	}

	for _, state := range t.States {
		for _, direction := range CardinalDirections {
			xNew, yNew := state.X+direction[0], state.Y+direction[1]

			if xNew < 0 || xNew >= len(board) || yNew < 0 || yNew >= len(board[xNew]) {
				continue
			}

			if state.Visited(xNew, yNew) {
				continue
			}

			if ch != board[xNew][yNew] {
				continue
			}

			child.States = append(child.States, &State{
				X:        xNew,
				Y:        yNew,
				Previous: state,
			})
		}
	}

	t.Children[ch] = child
	return child
}

func (t *Trie) FindWord(board [][]byte, word string) bool {
	if len(word) == 0 {
		return len(t.States) > 0
	}

	child, ok := t.Children[word[0]]
	if !ok {
		child = t.Explore(board, word[0])
	}

	return child.FindWord(board, word[1:])
}

func FindWords(board [][]byte, words []string) []string {
	trie := NewTrie(board)
	found := []string{}

	for _, word := range words {
		if trie.FindWord(board, word) {
			found = append(found, word)
		}
	}

	return found
}