package soko

import (
	"sort"
)

type PathFinder struct {
	state *State
	moves []Moves
}

func NewPathFinder(s *State) PathFinder {
	return PathFinder{s, nil}
}

func NewPathFinderMoves(s *State, moves []Moves) PathFinder {
	return PathFinder{s, moves}
}

func (p PathFinder) Find(target int) []int {
	source := p.state.Player.Pos
	level := p.state.Level
	size := level.Size()
	distances := NewInts(size)
	distances[source] = 0
	moves := NewInts(size)
	moves[source] = source
	state := p.newPathFinderState(source)

	heuristic := func(i int) int {
		idx := state.frontier[i]
		return distances[idx] + p.state.IdxToPos[idx].DistInt(p.state.IdxToPos[target])
	}
	for {
		curr := state.frontier[0]
		if curr == target {
			break
		}
		state.frontier = state.frontier[1:]

		state.findBetterNeighbours(distances, curr, func(nextIdx, newDistance int) {
			moves[nextIdx] = curr
		})

		if len(state.frontier) == 0 {
			return nil // no path
		}

		// apply heuristic to frontier (favor points closer to target)
		sort.Slice(state.frontier, func(i, j int) bool { return heuristic(i) < heuristic(j) })
	}

	// build reverse path
	curr := target
	path := []int{curr}
	for {
		curr = moves[curr]
		if curr == source {
			break
		}
		path = append(path, curr)
	}
	// reverse path
	n := len(path)
	for i := 0; i < n/2; i++ {
		path[i], path[n-i-1] = path[n-i-1], path[i]
	}
	return path
}

func (p PathFinder) FindDistances() Ints {
	source := p.state.Player.Pos
	level := p.state.Level
	size := level.Size()
	distances := NewInts(size)
	distances[source] = 0
	state := p.newPathFinderState(source)

	for {
		curr := state.frontier[0]
		state.frontier = state.frontier[1:]

		state.findBetterNeighbours(distances, curr, nil)

		if len(state.frontier) == 0 {
			return distances
		}
	}
}

func (p PathFinder) newPathFinderState(source int) *pathFinderState {
	return &pathFinderState{p.state, append(make([]int, 0, p.state.Level.Size()), source), p.moves}
}

type pathFinderState struct {
	state    *State
	frontier []int
	moves    []Moves
}

type betterNeighbourFn func(int, int)

func (s *pathFinderState) findBetterNeighbours(distances []int, curr int, better betterNeighbourFn) {
	currDistance := distances[curr]
	newDistance := currDistance + 1

	var moves []int
	if s.moves == nil {
		for _, dir := range Directions {
			nextIdx := s.state.Level.MoveIdx(curr, dir)
			if nextIdx == -1 {
				continue
			}
			moves = append(moves, nextIdx)
		}
	} else {
		moves = s.moves[curr].Valid
	}
	for _, nextIdx := range moves {
		if distance := distances[nextIdx]; distance != -1 && distance <= newDistance { // skip when shorter path exists
			continue
		}
		if !s.state.Walkable[nextIdx] || s.state.Bricks[nextIdx] != -1 { // filter neighbours
			continue
		}
		distances[nextIdx] = newDistance
		s.frontier = append(s.frontier, nextIdx)
		if better == nil {
			continue
		}
		better(nextIdx, newDistance)
	}
}