package n_queens

import (
	"bytes"
)

type board [][]bool

func newBoard(size int) board {
	b := make([][]bool, size)
	for i := range b {
		b[i] = make([]bool, size)
	}

	return board(b)
}

func (b board) size() (int, int) {
	return len(b[0]), len(b)
}

func (b board) get(x, y int) bool {
	return b[y][x]
}

func (b board) has(x, y int) bool {
	width, height := b.size()
	return x >= 0 && y >= 0 && x < width && y < height
}

func (b board) set(x, y int, value bool) {
	b[y][x] = value
}

type position struct{ x, y int }

type step struct {
	position int
	blocked  []position
}

type state struct {
	board board
	steps []step
	n     int
}

func newState(n int) *state {
	return &state{
		board: newBoard(n),
		n:     n,
		steps: make([]step, 0, n),
	}
}

func (s *state) options() []int {
	y := len(s.steps)
	if y >= s.n {
		return []int{}
	}

	results := []int{}
	for x := range s.n {
		if s.board.get(x, y) {
			continue
		}

		results = append(results, x)
	}

	return results
}

func (s *state) forward(x int) {
	blocked := []position{}
	y := len(s.steps)

	// Vertical
	for i := range s.n {
		if !s.board.get(x, i) {
			s.board.set(x, i, true)
			blocked = append(blocked, position{x, i})
		}
	}

	// Horizontal
	for i := range s.n {
		if !s.board.get(i, y) {
			s.board.set(i, y, true)
			blocked = append(blocked, position{i, y})
		}
	}

	// Forward vertical
	for offset := range s.n {
		p := position{offset, y + x - offset}

		if s.board.has(p.x, p.y) && !s.board.get(p.x, p.y) {
			s.board.set(p.x, p.y, true)
			blocked = append(blocked, p)
		}
	}

	// Backward vertical
	for offset := range s.n {
		p := position{offset, y - x + offset}

		if s.board.has(p.x, p.y) && !s.board.get(p.x, p.y) {
			s.board.set(p.x, p.y, true)
			blocked = append(blocked, p)
		}
	}

	s.steps = append(s.steps, step{
		position: x,
		blocked:  blocked,
	})
}

func (s *state) backward() {
	top, rest := s.steps[len(s.steps)-1], s.steps[:len(s.steps)-1]
	s.steps = rest

	for _, p := range top.blocked {
		s.board.set(p.x, p.y, false)
	}
}

func (s *state) solution() ([]string, bool) {
	if len(s.steps) != s.n {
		return nil, false
	}

	blank := bytes.Repeat([]byte{'.'}, s.n)

	solution := make([]string, s.n)
	for i := range solution {
		row := bytes.Clone(blank)
		row[s.steps[i].position] = 'Q'
		solution[i] = string(row)
	}

	return solution, true
}

func SolveNQueens(n int) [][]string {
	s := newState(n)
	solutions := [][]string{}

	var dfs func()
	dfs = func() {
		if soln, ok := s.solution(); ok {
			solutions = append(solutions, soln)
			return
		}

		for _, opt := range s.options() {
			s.forward(opt)
			dfs()
			s.backward()
		}
	}

	dfs()

	return solutions
}