package day16

import (
	"adventofcode2024/utils/dijkstra"
	"adventofcode2024/utils/grid2d"
	"adventofcode2024/utils/inputs"
	"fmt"
	_ "strings"
)

type (
	Vec  struct{ x, y int }
	Dir  struct{ x, y int }
	Cell rune
	Move struct {
		p Vec
		d Dir
	}
	Maze struct {
		height, width int
		start         Vec
		end           Vec
		grid          *grid2d.Grid[Cell]
	}
)

const (
	WALL     Cell = '#'
	CORRIDOR Cell = '.'
	START    Cell = 'S'
	END      Cell = 'E'
	PATH     Cell = '*'
)

var (
	NORTH      = Dir{0, -1}
	EAST       = Dir{1, 0}
	SOUTH      = Dir{0, 1}
	WEST       = Dir{-1, 0}
	DIRECTIONS = []Dir{NORTH, EAST, SOUTH, WEST}
)

func Part1(input string) int {
	maze := Maze{}.Parse(input)
	fmt.Println(maze)
	data := inputGraph(*maze)
	graph := dijkstra.CreateGraph(data)
	data.From.Label = "EAST"
	path, dist := dijkstra.GetShortestPath(data.From, data.To, graph)
	for _, p := range path {
		maze.grid.Set(p.X, p.Y, PATH)
	}
	fmt.Println(maze)
	return dist
}
func Part2(input string) int {
	maze := Maze{}.Parse(input)
	fmt.Println(maze)
	data := inputGraph(*maze)
	graph := dijkstra.CreateGraph(data)
	data.From.Label = "START"
	data.To.Label = "NORTH"
	paths, _ := dijkstra.GetAllShortestPaths(data.From, data.To, graph)
	bestPaths := make(map[Vec]struct{})
	for _, path := range paths {
		for _, node := range path {
			bestPaths[Vec{x: node.X, y: node.Y}] = struct{}{}
		}
	}
	return len(bestPaths)
}

func inputGraph(maze Maze) dijkstra.InputGraph {
	data := dijkstra.InputGraph{}
	data.From = dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "START"}
	data.To = dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: "END"}
	for y := 0; y < maze.height; y++ {
		for x := 0; x < maze.width; x++ {
			if maze.grid.Get(x, y) == WALL {
				continue
			}
			for _, dir := range DIRECTIONS {
				nx, ny := x+dir.x, y+dir.y
				if nx < 0 || ny < 0 || nx >= maze.width || ny >= maze.height {
					continue
				}
				if maze.grid.Get(nx, ny) == WALL {
					continue
				}
				for _, d := range getAllowedDirections(dir) {
					ss := dirToString(dir)
					ds := dirToString(d)
					data.Graph = append(data.Graph, dijkstra.InputData{
						Source:      dijkstra.Point{X: x, Y: y, Label: ss},
						Destination: dijkstra.Point{X: nx, Y: ny, Label: ds},
						Weight: func() int {
							if dir == d {
								return 1
							} else {
								return 1001
							}
						}(),
					})
				}
			}
		}
	}

	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: "END"},
		Destination: dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: dirToString(EAST)},
		Weight:      0,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: "END"},
		Destination: dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: dirToString(WEST)},
		Weight:      0,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: "END"},
		Destination: dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: dirToString(NORTH)},
		Weight:      0,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: "END"},
		Destination: dijkstra.Point{X: maze.end.x, Y: maze.end.y, Label: dirToString(SOUTH)},
		Weight:      0,
	})

	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "START"},
		Destination: dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: dirToString(EAST)},
		Weight:      0,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "START"},
		Destination: dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: dirToString(WEST)},
		Weight:      2000,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "START"},
		Destination: dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: dirToString(NORTH)},
		Weight:      1000,
	})
	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "START"},
		Destination: dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: dirToString(SOUTH)},
		Weight:      1000,
	})

	data.Graph = append(data.Graph, dijkstra.InputData{
		Source:      dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "EAST"},
		Destination: dijkstra.Point{X: maze.start.x, Y: maze.start.y, Label: "NORTH"},
		Weight:      1000,
	})

	return data
}

func dirToString(dir Dir) string {
	switch dir {
	case NORTH:
		return "NORTH"
	case EAST:
		return "EAST"
	case SOUTH:
		return "SOUTH"
	case WEST:
		return "WEST"
	default:
		return "UNKNOWN"
	}
}

func getAllowedDirections(direction Dir) []Dir {
	switch direction {
	case NORTH:
		return []Dir{NORTH, EAST, WEST}
	case SOUTH:
		return []Dir{SOUTH, EAST, WEST}
	case WEST:
		return []Dir{WEST, NORTH, SOUTH}
	case EAST:
		return []Dir{EAST, NORTH, SOUTH}
	}
	return []Dir{}
}

func (m Maze) Parse(input string) *Maze {
	m.grid = inputs.ToGrid2D(input, "\n", "", '?', func(c string) Cell { return Cell(c[0]) })
	m.height, m.width = m.grid.SizeY(), m.grid.SizeX()
	for y := 0; y < m.height; y++ {
		for x := 0; x < m.width; x++ {
			c := m.grid.Get(x, y)
			switch c {
			case WALL, CORRIDOR:
				m.grid.Set(x, y, c)
			case START:
				m.grid.Set(x, y, CORRIDOR)
				m.start = Vec{x, y}
			case END:
				m.grid.Set(x, y, CORRIDOR)
				m.end = Vec{x, y}
			}
		}
	}
	return &m
}

func (m *Maze) String() string {
	s := ""
	for y := 0; y < m.height; y++ {
		for x := 0; x < m.width; x++ {
			if y == m.start.y && x == m.start.x {
				s += string(START)
			} else if y == m.end.y && x == m.end.x {
				s += string(END)
			} else {
				s += string(m.grid.Get(x, y))
			}
		}
		s += "\n"
	}
	return s
}