package main

import (
	"fmt"
	"runtime"
	"strings"
	"sync"

	"git.mstar.dev/mstar/aoc24/util"
	"git.mstar.dev/mstar/goutils/sliceutils"
)

type MapState int
type Map [][]MapState // Y, X

type State struct {
	X, Y int
	Dir  MapState
}

const (
	MapStateEmpty MapState = iota
	MapStateObstacle
	MapStateVisited
	MapStateGuardUp
	MapStateGuardDown
	MapStateGuardLeft
	MapStateGuardRight
)

var amountOfGoroutines = 1

func parseStringIntoMapLine(line string) []MapState {
	mapLine := []MapState{}
	for _, r := range line {
		switch r {
		case '.':
			mapLine = append(mapLine, MapStateEmpty)
		case '#':
			mapLine = append(mapLine, MapStateObstacle)
		case '^':
			mapLine = append(mapLine, MapStateGuardUp)
		case '>':
			mapLine = append(mapLine, MapStateGuardRight)
		case '<':
			mapLine = append(mapLine, MapStateGuardLeft)
		case 'v':
			mapLine = append(mapLine, MapStateGuardDown)
		default:
		}
	}
	return mapLine
}

func parseLinesIntoMap(lines []string) Map {
	outMap := Map{}
	for _, line := range lines {
		outMap = append(outMap, parseStringIntoMapLine(line))
	}
	return outMap
}

func (area Map) isGuardInMap() bool {
	for _, areaLine := range area {
		for _, cell := range areaLine {
			switch cell {
			case MapStateGuardUp, MapStateGuardDown, MapStateGuardLeft, MapStateGuardRight:
				return true
			}
		}
	}
	return false
}

func (area Map) getGuardPosition() (x, y int, ok bool) {
	for iy, line := range area {
		for ix, cell := range line {
			switch cell {
			case MapStateGuardUp, MapStateGuardDown, MapStateGuardLeft, MapStateGuardRight:
				return ix, iy, true
			}
		}
	}
	return -1, -1, false
}

func (area Map) advanceMap() State {
	// Find guard position first
	guardX, guardY, ok := area.getGuardPosition()
	if !ok {
		return State{}
	}
	dir := area[guardY][guardX]
	switch area[guardY][guardX] {
	case MapStateGuardDown:
		targetCellVal := area.get(guardX, guardY+1)
		if targetCellVal == MapStateObstacle {
			area.set(MapStateGuardLeft, guardX, guardY)
		} else {
			area.set(MapStateVisited, guardX, guardY)
			area.set(MapStateGuardDown, guardX, guardY+1)
		}
	case MapStateGuardUp:
		targetCellVal := area.get(guardX, guardY-1)
		if targetCellVal == MapStateObstacle {
			area.set(MapStateGuardRight, guardX, guardY)
		} else {
			area.set(MapStateVisited, guardX, guardY)
			area.set(MapStateGuardUp, guardX, guardY-1)

		}
	case MapStateGuardLeft:
		targetCellVal := area.get(guardX-1, guardY)
		if targetCellVal == MapStateObstacle {
			area.set(MapStateGuardUp, guardX, guardY)
		} else {
			area.set(MapStateVisited, guardX, guardY)
			area.set(MapStateGuardLeft, guardX-1, guardY)
		}
	case MapStateGuardRight:
		targetCellVal := area.get(guardX+1, guardY)
		if targetCellVal == MapStateObstacle {
			area.set(MapStateGuardDown, guardX, guardY)
		} else {
			area.set(MapStateVisited, guardX, guardY)
			area.set(MapStateGuardRight, guardX+1, guardY)
		}
	default:
		panic(fmt.Sprintf("Unknown guard state %v", area[guardY][guardX]))
	}
	return State{guardX, guardY, dir}
}

func (area Map) set(newState MapState, posX, posY int) {
	if posY >= len(area) || posY < 0 || posX >= len(area[0]) || posX < 0 {
		return
	}
	area[posY][posX] = newState
}

func (area Map) get(posX, posY int) MapState {
	if posY >= len(area) || posY < 0 || posX >= len(area[0]) || posX < 0 {
		return MapStateEmpty
	}
	return area[posY][posX]
}

func (area Map) countVisitedCells() int {
	acc := 0
	for _, line := range area {
		for _, elem := range line {
			if elem == MapStateVisited {
				acc += 1
			}
		}
	}
	return acc
}

func (area Map) printCurrentState() string {
	builder := strings.Builder{}
	for _, line := range area {
		for _, elem := range line {
			switch elem {
			case MapStateEmpty:
				builder.WriteString(".")
			case MapStateObstacle:
				builder.WriteString("#")
			case MapStateVisited:
				builder.WriteString("X")
			case MapStateGuardDown:
				builder.WriteString("v")
			case MapStateGuardUp:
				builder.WriteString("^")
			case MapStateGuardLeft:
				builder.WriteString("<")
			case MapStateGuardRight:
				builder.WriteString(">")
			}
		}
		builder.WriteString("\n")
	}
	return builder.String()
}

func (area Map) countEmptyCells() uint64 {
	var acc uint64 = 0
	for _, line := range area {
		for _, elem := range line {
			if elem == MapStateEmpty {
				acc++
			}
		}
	}
	return acc
}

func (area Map) getFieldGuardIsFacing() (MapState, int, int) {
	guardX, GuardY, found := area.getGuardPosition()
	// No guard on map, not what we want
	if !found {
		return -1, -1, -1
	}
	guard := area.get(guardX, GuardY)
	switch guard {
	case MapStateGuardLeft:
		return area.get(guardX-1, GuardY), guardX - 1, GuardY
	case MapStateGuardRight:
		return area.get(guardX+1, GuardY), guardX + 1, GuardY
	case MapStateGuardDown:
		return area.get(guardX, GuardY+1), guardX, GuardY + 1
	case MapStateGuardUp:
		return area.get(guardX, GuardY-1), guardX, GuardY - 1
	default:
		panic(fmt.Errorf("Invalid field found by guard finder: %v", guard))
	}
}

func (area Map) clone() Map {
	clone := Map{}
	for _, line := range area {
		newLine := []MapState{}
		for _, elem := range line {
			newLine = append(newLine, elem)
		}
		clone = append(clone, newLine)
	}
	return clone
}

func runCheck(area Map, addChan chan any, wg *sync.WaitGroup, isRoot bool) {
	checkedStates := []State{}
	for area.isGuardInMap() {
		if isRoot {
			activeRoutines := runtime.NumGoroutine()
			if activeRoutines > amountOfGoroutines {
				amountOfGoroutines = activeRoutines
			}
			nextElem, nextX, nextY := area.getFieldGuardIsFacing()
			// fmt.Printf("Guard in root is facing %d:%d of type %d\n", nextX, nextY, nextElem)
			// If is root and next field is empty, launch a new goroutine with that field being an obstacle
			if nextElem == MapStateEmpty {
				clone := area.clone()
				clone.set(MapStateObstacle, nextX, nextY)
				// fmt.Printf(
				// 	"Launching check for new state. State:\n%s\n\n",
				// 	clone.printCurrentState(),
				// )
				fmt.Println("Launching new alternative from root")
				wg.Add(1)
				go runCheck(clone, addChan, wg, false)
			}
		}
		nextState := area.advanceMap()
		// Check if state is already known
		if len(
			sliceutils.Filter(
				checkedStates,
				func(t State) bool { return t.Dir == nextState.Dir && t.X == nextState.X && t.Y == nextState.Y },
			),
		) > 0 {
			// If it is, loop detected, report it back and exit
			wg.Done()
			addChan <- true
			fmt.Printf("Loop found, returning. IsRoot: %v\n", isRoot)
			return
		} else {
			checkedStates = append(checkedStates, nextState)
		}
	}
	// Guard no longer on map, report as done without loop
	wg.Done()
	fmt.Printf("Guard exited, returning. IsRoot: %v\n", isRoot)
}

func main() {
	inputLines := util.FileContentToNonEmptyLines(util.LoadFileFromArgs())
	// fullInput := sliceutils.Compact(
	// 	inputLines,
	// 	func(acc, next string) string { return acc + "\n" + next },
	// )
	area := parseLinesIntoMap(inputLines)
	for area.isGuardInMap() {
		area.advanceMap()
		// area.printCurrentState()
	}
	fmt.Printf("Task 1: %d\n", area.countVisitedCells())

	untouchedArea := parseLinesIntoMap(inputLines)

	acc := 0
	addChan := make(chan any, 1)
	wg := sync.WaitGroup{}
	go func() {
		for range addChan {
			acc++
		}
	}()
	wg.Add(1)
	go runCheck(untouchedArea.clone(), addChan, &wg, true)
	wg.Wait()
	close(addChan)
	fmt.Printf("Task 2: %d\n", acc)
	fmt.Printf("Peak nr of goroutines: %d\n", amountOfGoroutines)
}