package main

import (
	"fmt"
	"math"
	"slices"

	"git.mstar.dev/mstar/aoc24/util"
	"git.mstar.dev/mstar/goutils/sliceutils"
	"github.com/RyanCarrier/dijkstra/v2"
)

type Crossing struct {
	At    util.Vec2
	Exits []util.Vec2
}

type Connection struct {
	From, To util.Vec2
	Distance int
}

var (
	DirUp    = util.Vec2{X: 0, Y: -1}
	DirRight = util.Vec2{X: 1, Y: 0}
	DirDown  = util.Vec2{X: 0, Y: 1}
	DirLeft  = util.Vec2{X: -1, Y: 0}
)

// Get the size of the bord
func getSize(lines []string) util.Vec2 {
	return util.Vec2{X: int64(len(lines[0])), Y: int64(len(lines))}
}

// Count the paths next to a given positon
func countPathsAt(lines [][]rune, at util.Vec2) []util.Vec2 {
	center := lines[at.Y][at.X]
	paths := []util.Vec2{}
	if center == '#' {
		return paths
	}
	if util.SafeGet(lines, at.Up()) != '#' {
		paths = append(paths, at.Up())
	}
	if util.SafeGet(lines, at.Down()) != '#' {
		paths = append(paths, at.Down())
	}
	if util.SafeGet(lines, at.Left()) != '#' {
		paths = append(paths, at.Left())
	}
	if util.SafeGet(lines, at.Right()) != '#' {
		paths = append(paths, at.Right())
	}
	return paths
}

func sortPaths(area [][]rune) (all, straights, bends, crossings, deadEnds []util.Vec2) {
	for iy, line := range area {
		for ix, char := range line {
			if char == '#' {
				continue
			}
			pos := util.Vec2{X: int64(ix), Y: int64(iy)}
			all = append(all, pos)
			entries := countPathsAt(area, pos)
			switch len(entries) {
			case 0:
				continue
			case 1:
				deadEnds = append(deadEnds, pos)
			case 2:
				if entries[0].Right().Right().Eq(entries[1]) ||
					entries[0].Left().Left().Eq(entries[1]) ||
					entries[0].Up().Up().Eq(entries[1]) ||
					entries[0].Down().Down().Eq(entries[1]) {
					straights = append(straights, pos)
				} else {
					bends = append(bends, pos)
				}
			case 3, 4:
				crossings = append(crossings, pos)
			default:
				continue
			}
		}
	}
	return
}

func findAdjacentTargets(
	area [][]rune,
	at util.Vec2,
	crossings, bends, deadEnds []util.Vec2,
) []util.Vec2 {
	crossings = sliceutils.Filter(crossings, func(v util.Vec2) bool { return !v.Eq(at) })
	bends = sliceutils.Filter(bends, func(v util.Vec2) bool { return !v.Eq(at) })
	deadEnds = sliceutils.Filter(deadEnds, func(v util.Vec2) bool { return !v.Eq(at) })
	var posX, posY, negX, negY *util.Vec2
	if util.SafeGet(area, at.Right()) != '#' {
		crossingsPosX := sliceutils.Filter(
			crossings,
			func(t util.Vec2) bool { return at.X < t.X && at.Y == t.Y },
		)
		slices.SortFunc(crossingsPosX, func(a, b util.Vec2) int { return int(a.X - b.X) })
		if len(crossingsPosX) > 0 {
			posX = &crossingsPosX[0]
		}

		bendsPosX := sliceutils.Filter(
			bends,
			func(t util.Vec2) bool { return at.X < t.X && at.Y == t.Y },
		)
		slices.SortFunc(bendsPosX, func(a, b util.Vec2) int { return int(a.X - b.X) })
		if len(bendsPosX) > 0 {
			if posX == nil {
				posX = &bendsPosX[0]
			} else {
				if util.AbsI(posX.X-at.X) > util.AbsI(bendsPosX[0].X-at.X) {
					posX = &bendsPosX[0]
				}
			}
		}

		deadEndsPosX := sliceutils.Filter(
			deadEnds,
			func(t util.Vec2) bool { return at.X < t.X && at.Y == t.Y },
		)
		slices.SortFunc(deadEndsPosX, func(a, b util.Vec2) int { return int(a.X - b.X) })
		if len(deadEndsPosX) > 0 {
			if posX == nil {
				posX = &deadEndsPosX[0]
			} else {
				if util.AbsI(posX.X-at.X) > util.AbsI(deadEndsPosX[0].X-at.X) {
					posX = &deadEndsPosX[0]
				}
			}
		}
	}
	if util.SafeGet(area, at.Left()) != '#' {
		crossingsNegX := sliceutils.Filter(
			crossings,
			func(t util.Vec2) bool { return at.X > t.X && at.Y == t.Y },
		)
		slices.SortFunc(crossingsNegX, func(a, b util.Vec2) int { return int(b.X - a.X) })
		if len(crossingsNegX) > 0 {
			negX = &crossingsNegX[0]
		}
		bendsNegX := sliceutils.Filter(
			bends,
			func(t util.Vec2) bool { return at.X > t.X && at.Y == t.Y },
		)
		slices.SortFunc(bendsNegX, func(a, b util.Vec2) int { return int(b.X - a.X) })
		if len(bendsNegX) > 0 {
			if negX == nil {
				negX = &bendsNegX[0]
			} else {
				if util.AbsI(negX.X-at.X) > util.AbsI(bendsNegX[0].X-at.X) {
					negX = &bendsNegX[0]
				}
			}
		}
		deadEndsNegX := sliceutils.Filter(
			deadEnds,
			func(t util.Vec2) bool { return at.X > t.X && at.Y == t.Y },
		)
		slices.SortFunc(deadEndsNegX, func(a, b util.Vec2) int { return int(b.X - a.X) })
		if len(deadEndsNegX) > 0 {
			if negX == nil {
				negX = &deadEndsNegX[0]
			} else {
				if util.AbsI(negX.X-at.X) > util.AbsI(deadEndsNegX[0].X-at.X) {
					negX = &deadEndsNegX[0]
				}
			}
		}
	}
	if util.SafeGet(area, at.Down()) != '#' {
		crossingsPosY := sliceutils.Filter(
			crossings,
			func(t util.Vec2) bool { return at.Y < t.Y && at.X == t.X },
		)
		slices.SortFunc(crossingsPosY, func(a, b util.Vec2) int { return int(a.Y - b.Y) })
		if len(crossingsPosY) > 0 {
			posY = &crossingsPosY[0]
		}
		bendsPosY := sliceutils.Filter(
			bends,
			func(t util.Vec2) bool { return at.Y < t.Y && at.X == t.X },
		)
		slices.SortFunc(bendsPosY, func(a, b util.Vec2) int { return int(a.Y - b.Y) })
		if len(bendsPosY) > 0 {
			if posY == nil {
				posY = &bendsPosY[0]
			} else {
				if util.AbsI(posY.Y-at.Y) < util.AbsI(bendsPosY[0].Y-at.Y) {
					posY = &bendsPosY[0]
				}
			}
		}
		deadEndsPosY := sliceutils.Filter(
			deadEnds,
			func(t util.Vec2) bool { return at.Y < t.Y && at.X == t.X },
		)
		slices.SortFunc(deadEndsPosY, func(a, b util.Vec2) int { return int(a.Y - b.Y) })
		if len(deadEndsPosY) > 0 {
			if posY == nil {
				posY = &deadEndsPosY[0]
			} else {
				if util.AbsI(posY.Y-at.Y) < util.AbsI(deadEndsPosY[0].Y-at.Y) {
					posY = &deadEndsPosY[0]
				}
			}
		}
	}
	if util.SafeGet(area, at.Up()) != '#' {
		crossingsNegY := sliceutils.Filter(
			crossings,
			func(t util.Vec2) bool { return at.Y > t.Y && at.X == t.X },
		)
		slices.SortFunc(crossingsNegY, func(a, b util.Vec2) int { return int(b.Y - a.Y) })
		if len(crossingsNegY) > 0 {
			negY = &crossingsNegY[0]
		}
		bendsNegY := sliceutils.Filter(
			bends,
			func(t util.Vec2) bool { return at.Y > t.Y && at.X == t.X },
		)
		slices.SortFunc(bendsNegY, func(a, b util.Vec2) int { return int(b.Y - a.Y) })
		if len(bendsNegY) > 0 {
			if negY == nil {
				negY = &bendsNegY[0]
			} else {
				if util.AbsI(negY.Y-at.Y) > util.AbsI(bendsNegY[0].Y-at.Y) {
					negY = &bendsNegY[0]
				}
			}
		}
		deadEndsNegY := sliceutils.Filter(
			deadEnds,
			func(t util.Vec2) bool { return at.Y > t.Y && at.X == t.X },
		)
		slices.SortFunc(deadEndsNegY, func(a, b util.Vec2) int { return int(b.Y - a.Y) })
		if len(deadEndsNegY) > 0 {
			if negY == nil {
				negY = &deadEndsNegY[0]
			} else {
				if util.AbsI(negY.Y-at.Y) > util.AbsI(deadEndsNegY[0].Y-at.Y) {
					negY = &deadEndsNegY[0]
				}
			}
		}
	}
	// fmt.Printf("posX %#v, negX %#v, posY %#v, negY %#v\n", posX, negX, posY, negY)
	combined := []util.Vec2{}
	if posX != nil {
		combined = append(combined, *posX)
	}
	if negX != nil {
		combined = append(combined, *negX)
	}
	if posY != nil {
		combined = append(combined, *posY)
	}
	if negY != nil {
		combined = append(combined, *negY)
	}
	return combined
}

// Get the distance between two directly connecting points which must be either a crossing or a dead end
func distanceBetween(
	area [][]rune,
	start, end util.Vec2,
	crossings, bends, deadEnds []util.Vec2,
) (distance int, ok bool) {
	// Locations are the same, note that and exit
	if start.Eq(end) {
		return 0, true
	}
	// TODO: Refactor slightly to allow for distance to straights and bends too

	// Guard against start or end not being either a crossing or dead end
	if !(sliceutils.Contains(crossings, start) || sliceutils.Contains(deadEnds, start)) &&
		!(sliceutils.Contains(crossings, end) || sliceutils.Contains(deadEnds, end)) {
		return 0, false
	}
	// Filter out start from crossing and dead ends to avoid recursion or zero length paths
	// crossings = sliceutils.Filter(crossings, func(v util.Vec2) bool { return !v.Eq(start) })
	// deadEnds = sliceutils.Filter(deadEnds, func(v util.Vec2) bool { return !v.Eq(start) })
	adjacents := findAdjacentTargets(area, start, crossings, bends, deadEnds)
	for _, adj := range adjacents {
		distanceToStart := int(math.Round(start.Add(adj.Mult(-1)).Len()))
		trackEnd, trackLen := followBend(area, adj, start, crossings, bends, deadEnds)
		if trackEnd.Eq(end) {
			return distanceToStart + trackLen, true
		}
	}
	return -1, false
}

func followBend(
	area [][]rune,
	at, from util.Vec2,
	crossings, bends, deadEnds []util.Vec2,
) (nextCrossingOrDeadEnd util.Vec2, distanceTravelled int) {
	if sliceutils.Contains(crossings, at) {
		return at, 0
	}
	if sliceutils.Contains(deadEnds, at) {
		return at, 0
	}
	adjacents := findAdjacentTargets(area, at, crossings, bends, deadEnds)
	next := sliceutils.Filter(adjacents, func(v util.Vec2) bool { return !v.Eq(from) })[0]
	diff := at.Add(next.Mult(-1))
	diffLen := uint64(math.Round(diff.Len()))
	end, lenToEnd := followBend(area, next, at, crossings, bends, deadEnds)

	return end, int(diffLen) + 1000 + lenToEnd
}

func findAdjacentCrossingsAndDeadEnds(
	area [][]rune,
	at util.Vec2,
	crossings, bends, deadEnds []util.Vec2,
) []util.Vec2 {
	adjs := findAdjacentTargets(area, at, crossings, bends, deadEnds)
	// fmt.Println("adjacents", adjs)
	return sliceutils.Map(adjs, func(t util.Vec2) util.Vec2 {
		end, _ := followBend(area, t, at, crossings, bends, deadEnds)
		return end
	})
}

func buildGraph(
	area [][]rune,
	crossings, bends, deadEnds []util.Vec2,
	start, end util.Vec2,
) (graph *dijkstra.Graph, poi []util.Vec2, err error) {
	g := dijkstra.NewGraph()
	graph = &g
	poi = slices.Clone(crossings)
	poi = append(poi, deadEnds...)
	poi = append(poi, start, end)
	// Init all vertecies first
	for i := range poi {
		err = graph.AddEmptyVertex(i)
		if err != nil {
			return
		}
	}
	for i, pos := range poi {
		adjs := findAdjacentCrossingsAndDeadEnds(area, pos, crossings, bends, deadEnds)
		for _, adj := range adjs {
			index := slices.Index(poi, adj)
			if index == -1 {
				panic("Didn't find point in poi")
			}
			distance, _ := distanceBetween(area, pos, adj, crossings, bends, deadEnds)
			err = graph.AddArc(i, index, uint64(distance))
			if err != nil {
				return
			}
		}
	}

	// TODO: Manually add end inbound connections
	for _, locs := range findAdjacentCrossingsAndDeadEnds(area, end, crossings, bends, deadEnds) {
		_ = locs
	}
	return
}

func findStartEnd(area [][]rune) (start util.Vec2, end util.Vec2) {
	for iy, line := range area {
		for ix, char := range line {
			if char == 'S' {
				start = util.Vec2{X: int64(ix), Y: int64(iy)}
			} else if char == 'E' {
				end = util.Vec2{X: int64(ix), Y: int64(iy)}
			}
		}
	}
	return
}

func main() {
	inputLines := util.FileContentToNonEmptyLines(util.LoadFileFromArgs())
	// size := getSize(inputLines)
	area := sliceutils.Map(inputLines, func(t string) []rune { return []rune(t) })
	start, end := findStartEnd(area)
	_, _, bends, crossings, deadEnds := sortPaths(area)
	// fmt.Println("straights:", straights)
	// fmt.Println("bends:", bends)
	// fmt.Println("crossings:", crossings)
	// fmt.Println("dead ends:", deadEnds)
	fmt.Println(
		"start adjacents:",
		findAdjacentCrossingsAndDeadEnds(area, start, crossings, bends, deadEnds),
	)
	fmt.Println(
		"end adjacents:",
		findAdjacentCrossingsAndDeadEnds(area, util.Vec2{X: 9, Y: 7}, crossings, bends, deadEnds),
	)
	dist, ok := distanceBetween(area, util.Vec2{X: 9, Y: 7}, end, crossings, bends, deadEnds)
	fmt.Println("dist, ok:", dist, ok)
	// fmt.Println("dist", distanceBetween(area [][]rune, start util.Vec2, end util.Vec2, crossings []util.Vec2, bends []util.Vec2, deadEnds []util.Vec2))
	graph, pointsOfInterest, err := buildGraph(area, crossings, bends, deadEnds, start, end)
	if err != nil {
		panic(err)
	}
	shortest, err := graph.Shortest(
		slices.Index(pointsOfInterest, start),
		slices.Index(pointsOfInterest, end),
	)
	if err != nil {
		panic(err)
	}
	fmt.Println("shortest", shortest)
}