package main

import (
	"fmt"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := h.StdinToStringSlice()
	part1(inp)
	fmt.Println()
	part2(inp)
}

func part1(input []string) {
	m := h.StringSliceToCoordByteMap(input)

	fmt.Println("# Part 1")
	fmt.Println(solve(&m, 0, 3))
}

// 1188 is too low
func part2(input []string) {
	m := h.StringSliceToCoordByteMap(input)

	fmt.Println("# Part 2")
	fmt.Println(solve(&m, 4, 10))
}

func solve(m *h.CoordByteMap, minStraight, maxStraight int) int {
	q := h.NewHeap[heatState](func(a, b heatState) bool {
		return a.heatLoss < b.heatLoss
	})
	q.Push(heatState{
		state: state{
			pos:      &h.Coordinate{X: m.TLX + 1, Y: m.TLY},
			dir:      dirE,
			dirCount: 1,
		},
	})
	q.Push(heatState{
		state: state{
			pos:      &h.Coordinate{X: m.TLX, Y: m.TLY + 1},
			dir:      dirS,
			dirCount: 1,
		},
	})
	end := h.Coordinate{X: m.BRX, Y: m.BRY}
	visitedCoords := make(map[h.Coordinate]bool)
	visited := make(map[string]int)

	for q.Len() != 0 {
		chk := q.Pop()
		if !m.ContainsCoord(*chk.state.pos) {
			continue
		}
		heat := int(m.Get(*chk.state.pos)-'0') + chk.heatLoss
		if chk.state.pos.Equals(end) && chk.dirCount >= minStraight {
			return heat
		}
		if visHeat, ok := visited[chk.state.String()]; ok {
			if visHeat <= heat {
				continue
			}
		}
		visitedCoords[*chk.state.pos] = true
		visited[chk.state.String()] = heat
		if chk.dirCount >= minStraight {
			// Add Left Turn
			left := chk.dir.LeftTurn().Get(chk.pos)
			if m.ContainsCoord(*left) {
				q.Push(heatState{
					state: state{
						pos:      left,
						dir:      chk.dir.LeftTurn(),
						dirCount: 1,
					},
					heatLoss: heat,
				})
			}
			// Add Right Turn
			right := chk.dir.RightTurn().Get(chk.pos)
			if m.ContainsCoord(*right) {
				q.Push(heatState{
					state: state{
						pos:      right,
						dir:      chk.dir.RightTurn(),
						dirCount: 1,
					},
					heatLoss: heat,
				})
			}
		}

		if chk.dirCount < maxStraight {
			// Add Straight
			straight := chk.dir.Get(chk.pos)
			if m.ContainsCoord(*straight) {
				q.Push(heatState{
					state: state{
						pos:      straight,
						dir:      chk.dir,
						dirCount: chk.dirCount + 1,
					},
					heatLoss: heat,
				})
			}
		}
	}
	return h.MAX_INT
}

type state struct {
	pos      *h.Coordinate
	dir      direction
	dirCount int
}

func (s state) String() string {
	return fmt.Sprintf("%s-%s-%d", s.pos, s.dir, s.dirCount)
}

type heatState struct {
	state
	heatLoss int
}

type direction h.Coordinate

func (d direction) Get(from *h.Coordinate) *h.Coordinate {
	switch d.String() {
	case "N":
		return from.GetNorthCoord()
	case "E":
		return from.GetEastCoord()
	case "S":
		return from.GetSouthCoord()
	case "W":
		return from.GetWestCoord()
	}
	return from
}

func (d direction) LeftTurn() direction {
	switch d.String() {
	case "N":
		return dirW
	case "E":
		return dirN
	case "S":
		return dirE
	default:
		return dirS
	}
}

func (d direction) RightTurn() direction {
	switch d.String() {
	case "N":
		return dirE
	case "E":
		return dirS
	case "S":
		return dirW
	default:
		return dirN
	}
}

func (d direction) String() string {
	switch {
	case d.X == dirN.X && d.Y == dirN.Y:
		return "N"
	case d.X == dirE.X && d.Y == dirE.Y:
		return "E"
	case d.X == dirS.X && d.Y == dirS.Y:
		return "S"
	default:
		return "W"
	}
}

var (
	dirN = direction{X: 0, Y: -1}
	dirE = direction{X: 1, Y: 0}
	dirS = direction{X: 0, Y: 1}
	dirW = direction{X: -1, Y: 0}
)