package maze

import (
	"image"
	"image/color"
	"image/draw"
	"math/rand"
)

type cell uint8

const (
	// cellRight is set if movement to the right is allowed.
	cellRight cell = 1 << iota

	// cellDown is set if movement down is allowed.
	cellDown
)

type Maze struct {
	Size       image.Rectangle
	Start, End image.Point

	// the cells in the maze, starting at size.Min and going left-to-right, top-to-bottom.
	cells []cell
}

// if cells at i and j belong to different neighborhoods, they will be merged
// and true will be returned. Otherwise, false will be returned.
func tryJoinNeighborhoods(neighborhood []*[]int, i, j int) bool {
	if aN, bN := neighborhood[i], neighborhood[j]; aN != bN {
		// prefer the larger set absorbing the smaller set.
		if len(*aN) < len(*bN) {
			aN, bN = bN, aN
		}

		// aN is now the larger set, and bN is the smaller set.

		// append the smaller set to the larger set.
		*aN = append(*aN, *bN...)

		// update the membership of the smaller set to point to the larger set.
		for _, member := range *bN {
			neighborhood[member] = aN
		}

		return true
	}
	return false
}

// Make creates a new maze of the given size, using the given random number generator.
//
// If r is nil, the default random number generator will be used.
func Make(size image.Rectangle, r *rand.Rand) Maze {
	if size.Dx() <= 0 || size.Dy() <= 0 {
		panic("invalid maze size")
	}

	if r == nil {
		// create a new random number generator if
		// one wasn't provided.
		r = rand.New(rand.NewSource(rand.Int63()))
	}

	w, h := size.Dx(), size.Dy()
	neighborhoods := make([]*[]int, w*h)
	cells := make([]cell, w*h)

	// each cell begins initially disconnected, alone in its own set.
	for i := range cells {
		neighborhoods[i] = &[]int{i}
	}

	// create a list of all edges in the maze.
	// there are (w-1)*h horizontal edges, and w*(h-1) vertical edges.
	// the first (w-1)*h edges are horizontal, and the rest are vertical.
	edges := make([]int, (w-1)*h+w*(h-1))
	for i := range edges {
		edges[i] = i
	}

	// shuffle the list of edges, this will be the order in which we attempt to join cells.
	r.Shuffle(len(edges), func(i, j int) {
		edges[i], edges[j] = edges[j], edges[i]
	})

	// repeatedly join cells until all cells are in the same set.
	for _, edge := range edges {
		if edge < (w-1)*h {
			// join right
			// we need to convert the edge index into coordinates, and then back into a cell index.
			i := edge%(w-1) + edge/(w-1)*w
			if tryJoinNeighborhoods(neighborhoods, i, i+1) {
				cells[i] |= cellRight
			} else {
				// already joined, skip to the next edge.
				continue
			}
		} else {
			// join down
			// we can mostly use the edge index unmodified, we just need to subtract the length of the horizontal edges.
			i := edge - (w-1)*h
			if tryJoinNeighborhoods(neighborhoods, i, i+w) {
				cells[i] |= cellDown
			} else {
				// already joined, skip to the next edge.
				continue
			}
		}

		// if we reach this point, then two neighborhoods were joined. If all cells are in the same neighborhood,
		// then we can stop trying to merge the remaining edges.
		if len(*neighborhoods[0]) == w*h {
			break
		}
	}

	var start, end image.Point

	// choose a start and end point on opposite sides of the maze.
	switch r.Intn(2) {
	case 0: // left/right
		start = image.Pt(0, r.Intn(h)).Add(size.Min)
		end = image.Pt(w-1, r.Intn(h)).Add(size.Min)

	case 1: // top/bottom
		start = image.Pt(r.Intn(w), 0).Add(size.Min)
		end = image.Pt(r.Intn(w), h-1).Add(size.Min)
	}

	if r.Intn(2) == 0 {
		// swap start and end points.
		start, end = end, start
	}

	return Maze{
		Size:  size,
		Start: start,
		End:   end,
		cells: cells,
	}
}

func (m Maze) coordToIndex(p image.Point) int {
	return p.X - m.Size.Min.X + (p.Y-m.Size.Min.Y)*m.Size.Dx()
}

// Returns true if you can move to the right from the given point.
func (m Maze) Right(p image.Point) bool {
	return p.In(m.Size) && m.cells[m.coordToIndex(p)]&cellRight != 0
}

// Returns true if you can move down from the given point.
func (m Maze) Down(p image.Point) bool {
	return p.In(m.Size) && m.cells[m.coordToIndex(p)]&cellDown != 0
}

// Returns true if you can move to the left from the given point.
func (m Maze) Left(p image.Point) bool {
	return m.Right(p.Sub(image.Pt(1, 0)))
}

// Returns true if you can move up from the given point.
func (m Maze) Up(p image.Point) bool {
	return m.Down(p.Sub(image.Pt(0, 1)))
}

// Returns a drawing of the maze, mostly for debugging purposes.
// An optional path can be drawn on top of the maze.
func (m Maze) Draw(path []image.Point) *image.RGBA {
	const cellSize = 16
	const markerRadius = 4 // the radius of the start and end markers.
	const wallRadius = 2   // half the thickness of the wall lines.
	const pathRadius = 2   // half the thickness of the path line.

	backgroundColor := image.NewUniform(color.Black)
	wallColor := image.NewUniform(color.White)
	startColor := image.NewUniform(color.RGBA{0, 255, 0, 255})
	endColor := image.NewUniform(color.RGBA{255, 0, 0, 255})
	pathColor := image.NewUniform(color.RGBA{0, 0, 255, 255})

	leftWall := image.Rect(0, 0, 0, cellSize).Inset(-wallRadius)
	topWall := image.Rect(0, 0, cellSize, 0).Inset(-wallRadius)
	rightWall := leftWall.Add(image.Pt(cellSize, 0))
	bottomWall := topWall.Add(image.Pt(0, cellSize))

	marker := image.Rect(cellSize/2-markerRadius, cellSize/2-markerRadius, cellSize/2+markerRadius, cellSize/2+markerRadius)

	img := image.NewRGBA(image.Rect(0, 0, m.Size.Dx()*cellSize+wallRadius*2, m.Size.Dy()*cellSize+wallRadius*2))

	// fill the entire image with the background color.
	draw.Draw(img, img.Bounds(), backgroundColor, image.Point{}, draw.Src)

	for y := m.Size.Min.Y; y < m.Size.Max.Y; y++ {
		for x := m.Size.Min.X; x < m.Size.Max.X; x++ {

			p := image.Pt(x, y)
			cellP := p.Sub(m.Size.Min).Mul(cellSize).Add(image.Pt(wallRadius, wallRadius))

			draw := func(rect image.Rectangle, color *image.Uniform) {
				draw.Draw(img, rect.Add(cellP), color, image.Point{}, draw.Src)
			}

			if p == m.Start {
				draw(marker, startColor)
			}

			if p == m.End {
				draw(marker, endColor)
			}

			if !m.Right(p) {
				draw(rightWall, wallColor)
			}

			if !m.Down(p) {
				draw(bottomWall, wallColor)
			}

			if !m.Left(p) {
				draw(leftWall, wallColor)
			}

			if !m.Up(p) {
				draw(topWall, wallColor)
			}
		}
	}

	if len(path) > 1 {
		pathDrawOffset := image.Pt(cellSize/2+wallRadius, cellSize/2+wallRadius)

		draw := func(i, j int) {
			// fake drawing a line between the two points by treating them as points
			// of a rectangle, and drawing that rectangle instead.
			//
			// this only works because we're assuming there are no diagonal movements.
			rect := image.Rectangle{
				path[i].Sub(m.Size.Min).Mul(cellSize),
				path[j].Sub(m.Size.Min).Mul(cellSize),
			}.Canon().Inset(-pathRadius).Add(pathDrawOffset)

			draw.Draw(img, rect, pathColor, image.Point{}, draw.Src)
		}

		start := 0

		// cheating by assuming adjacent points in the path are adjacent cells,
		// and movements are only horizontal or vertical.
		for i := 1; i < len(path); i++ {
			if path[start].X == path[i].X || path[start].Y == path[i].Y {
				// while points are on the same horizontal or vertical line,
				// skip the intermediate points.
				continue
			}

			// all the points before this one were on the same line, so draw them
			// as a single line.
			draw(start, i-1)
			start = i - 1
		}

		// any remaining points we haven't drawn yet are also on the same line.
		draw(start, len(path)-1)
	}

	return img
}