package tsp

import (
	"errors"
	"fmt"

	"smariot.com/tsp/internal/solver"
	"smariot.com/tsp/internal/solver/problem"
)

// Problem represents a search problem.
type Problem[State comparable, Solution any] interface {
	// Require a problem to implement the problem.Problem interface.
	problem.Problem[State]

	// Appends the initial states to begin the search with.
	Initialize(out []State) ([]State, error)

	// Appends the next possible states to the given slice.
	//
	// Appending an already known state is allowed, if this happens its
	// relative order is assumed to have changed.
	//
	// A state should only be resubmitted if its cost has changed, otherwise
	// we'd potentially be stuck in a loop trying to explore the same state over and over.
	Next(seed State, out []State) ([]State, error)

	// Returns true if this node represents a complete solution.
	Solved(state State) bool

	// Converts the given solved state into a solution that can be returned by the Solve function.
	Finish(state State) (Solution, error)
}

// ErrNoSolution is returned by Solve when all states have been checked without finding a solution.
var ErrNoSolution = errors.New("no solution found")

// ErrBadCapacity is returned by Solve when the capacity is not positive.
var ErrBadCapacity = errors.New("capacity must be positive")

// Solves the given problem, returning the best state found.
//
// The capacity is the maximum number of states that can be stored in memory at any given time.
//
// If capacity is greater than 0, then the solver will maintain a max heap and discard states when it reaches capacity.
func Solve[P Problem[State, Solution], State comparable, Solution any](problem P, capacity int) (Solution, error) {
	if capacity < 0 {
		var zero Solution
		return zero, fmt.Errorf("%w, got %d", ErrBadCapacity, capacity)
	}

	solver := solver.New(problem, capacity)

	next, err := problem.Initialize(nil)

	// note that we push the states even in the case of an error,
	// as it's simplifies the cleanup process.
	for _, state := range next {
		solver.Push(state)
	}

	if err != nil {
		solver.Reset()
		var zero Solution
		return zero, err
	}

	for {
		state, ok := solver.Pop()

		if !ok {
			var zero Solution
			return zero, ErrNoSolution
		}

		if problem.Solved(state) {
			solver.Reset()
			solution, err := problem.Finish(state)
			problem.Discard(state)
			return solution, err
		}

		next, err = problem.Next(state, next[:0])

		// again, we're going to deal with the states we were given first
		// before we deal with the error.
		resubmitted := false
		for _, nextState := range next {
			if state == nextState {
				resubmitted = true
			}

			solver.Push(nextState)
		}

		// problem is allowed to resubmit the state. If it did, then it's either somewhere in the heap,
		// or it already discarded it due to being at capacity. In either case, we shouldn't discard it again.
		if !resubmitted {
			problem.Discard(state)
		}

		if err != nil {
			solver.Reset()
			var zero Solution
			return zero, err
		}
	}
}