package day20

import (
	"container/heap"
	"fmt"
	"math"
	"strings"
)

var directions = [][]int{
	{0, 1},  // Right
	{1, 0},  // Down
	{0, -1}, // Left
	{-1, 0}, // Up
}

type State struct {
	row, col, dir, score int
	path                 []string
	//Would rather use tuple as key but cant get to work
	costMap map[[2]int]int
}

type PriorityQueue []State

func Part1(input string) int {
	maze := parseInput(input)
	_, path, costMap := solveMaze(maze)
	cheats := findCheats(path, costMap)
	return cheats
}

func Part2(input string) int {
	maze := parseInput(input)
	_, path, costMap := solveMaze(maze)
	cheats := findCheatsPart2(path, costMap)
	return cheats
}

func parseInput(input string) []string {
	parts := strings.Split(input, "\n")
	maze := []string{}
	maze = append(maze, parts...)

	return maze
}

func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool {
	return pq[i].score < pq[j].score
}
func (pq PriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }

func (pq *PriorityQueue) Push(x interface{}) {
	*pq = append(*pq, x.(State))
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	*pq = old[0 : n-1]
	return item
}

func solveMaze(maze []string) (int, []string, map[[2]int]int) {
	rows := len(maze)
	cols := len(maze[0])

	// Locate start and end points
	var startRow, startCol, endRow, endCol int
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			if maze[i][j] == 'S' {
				startRow, startCol = i, j
			} else if maze[i][j] == 'E' {
				endRow, endCol = i, j
			}
		}
	}

	// Priority queue and visited set
	pq := &PriorityQueue{}
	heap.Init(pq)
	startPath := []string{fmt.Sprintf("(%d,%d)", startRow, startCol)}
	startCostMap := map[[2]int]int{
		{startRow, startCol}: 0,
	}
	heap.Push(pq, State{startRow, startCol, 0, 0, startPath, startCostMap})
	visited := make(map[[3]int]bool)

	for pq.Len() > 0 {
		cur := heap.Pop(pq).(State)

		// If reached the end point, return the score and path
		if cur.row == endRow && cur.col == endCol {
			return cur.score, cur.path, cur.costMap
		}

		// Mark as visited
		key := [3]int{cur.row, cur.col, cur.dir}
		if visited[key] {
			continue
		}
		visited[key] = true

		// Explore moves
		for i := 0; i < 4; i++ {
			// Compute new direction
			newDir := (cur.dir + i) % 4

			// Compute new position
			newRow := cur.row + directions[newDir][0]
			newCol := cur.col + directions[newDir][1]

			// Check bounds and wall
			if newRow >= 0 && newRow < rows && newCol >= 0 && newCol < cols && maze[newRow][newCol] != '#' {
				newCost := cur.score + 1
				if oldCost, exists := cur.costMap[[2]int{newRow, newCol}]; !exists || newCost < oldCost {
					// Update cost map
					newCostMap := make(map[[2]int]int)
					for k, v := range cur.costMap {
						newCostMap[k] = v
					}
					newCostMap[[2]int{newRow, newCol}] = newCost

					// Update path
					newPath := append([]string{}, cur.path...)
					newPath = append(newPath, fmt.Sprintf("(%d,%d)", newRow, newCol))

					heap.Push(pq, State{
						newRow,
						newCol,
						newDir,
						newCost,
						newPath,
						newCostMap,
					})
				}
			}
		}
	}

	return -1, nil, nil // No path found
}

func findCheats(path []string, costLookup map[[2]int]int) int {

	totalCheats := 0
	for index, position := range path {

		var row, col int
		//Should be a better way to do this
		fmt.Sscanf(position, "(%d,%d)", &row, &col)

		for _, dir := range directions {
			newRow := row + (dir[0] * 2)
			newCol := col + (dir[1] * 2)

			if cost, exists := costLookup[[2]int{newRow, newCol}]; exists {
				if cost > index && (cost-index-2) >= 100 {
					totalCheats++
				}
			}
		}
	}
	return totalCheats
}

func findCheatsPart2(path []string, costLookup map[[2]int]int) int {
	totalCheats := 0

	for index, position := range path {
		var row, col int
		//Should be a better way to do this
		fmt.Sscanf(position, "(%d,%d)", &row, &col)

		for x1 := row - 20; x1 <= row+20; x1++ {
			for y1 := col - 20 + int(math.Abs(float64(row-x1))); y1 <= col+20-int(math.Abs(float64(row-x1))); y1++ {
				if cost, exists := costLookup[[2]int{x1, y1}]; exists && cost != -1 {
					diff := cost - index - int(math.Abs(float64(row-x1))) - int(math.Abs(float64(col-y1)))
					if cost > index && diff >= 100 {
						totalCheats++
					}
				}
			}
		}
	}

	return totalCheats
}