package day16

import (
	"fmt"
	"sync"

	"gitea.paas.celticinfo.fr/oabrivard/aoc2023/utils"
)

type Move struct {
	row       int
	col       int
	direction rune
}

func nextMoves(origin *Move, lines []string) []*Move {
	result := []*Move{}

	switch lines[origin.row][origin.col] {
	case '/':
		switch origin.direction {
		case 'R':
			result = append(result, &Move{origin.row - 1, origin.col, 'U'})
		case 'L':
			result = append(result, &Move{origin.row + 1, origin.col, 'D'})
		case 'U':
			result = append(result, &Move{origin.row, origin.col + 1, 'R'})
		case 'D':
			result = append(result, &Move{origin.row, origin.col - 1, 'L'})
		}
	case '\\':
		switch origin.direction {
		case 'R':
			result = append(result, &Move{origin.row + 1, origin.col, 'D'})
		case 'L':
			result = append(result, &Move{origin.row - 1, origin.col, 'U'})
		case 'U':
			result = append(result, &Move{origin.row, origin.col - 1, 'L'})
		case 'D':
			result = append(result, &Move{origin.row, origin.col + 1, 'R'})
		}
	case '|':
		switch origin.direction {
		case 'R', 'L':
			result = append(result, &Move{origin.row - 1, origin.col, 'U'})
			result = append(result, &Move{origin.row + 1, origin.col, 'D'})
		case 'U':
			result = append(result, &Move{origin.row - 1, origin.col, 'U'})
		case 'D':
			result = append(result, &Move{origin.row + 1, origin.col, 'D'})
		}
	case '-':
		switch origin.direction {
		case 'R':
			result = append(result, &Move{origin.row, origin.col + 1, 'R'})
		case 'L':
			result = append(result, &Move{origin.row, origin.col - 1, 'L'})
		case 'U', 'D':
			result = append(result, &Move{origin.row, origin.col - 1, 'L'})
			result = append(result, &Move{origin.row, origin.col + 1, 'R'})
		}
	default:
		switch origin.direction {
		case 'R':
			result = append(result, &Move{origin.row, origin.col + 1, 'R'})
		case 'L':
			result = append(result, &Move{origin.row, origin.col - 1, 'L'})
		case 'U':
			result = append(result, &Move{origin.row - 1, origin.col, 'U'})
		case 'D':
			result = append(result, &Move{origin.row + 1, origin.col, 'D'})
		}
	}

	return result
}

func Part1(lines []string, startRow int, startCol int, direction rune) int {
	height := len(lines)
	width := len(lines[0])
	cache := map[Move]bool{}
	isEnergized := make([][]bool, height)
	beamMoves := utils.NewStack[*Move]()

	for row := range lines {
		isEnergized[row] = make([]bool, width)
	}

	beamMoves.Push(&Move{startRow, startCol, direction})
	i := 0
	for beamMoves.HasElement() {
		currOrigin := beamMoves.Pop()
		isEnergized[currOrigin.row][currOrigin.col] = true

		_, hasKey := cache[*currOrigin]
		if hasKey {
			continue
		}

		cache[*currOrigin] = true
		moves := nextMoves(currOrigin, lines)

		for _, move := range moves {
			if move.row >= 0 && move.row < height && move.col >= 0 && move.col < width {
				beamMoves.Push(move)
			}
		}

		i++
		if i > 1000000 {
			fmt.Println("loop ?")
			return -1
		}
	}

	count := 0
	for row := range isEnergized {
		for col := range isEnergized[row] {
			if isEnergized[row][col] {
				count++
			}
		}
	}

	return count
}

func Part2(lines []string) int {
	height := len(lines)
	width := len(lines[0])
	max := -1

	var wgD sync.WaitGroup
	chD := make(chan int, width)

	for i := 0; i < width; i++ {
		wgD.Add(1)
		go func(col int) {
			defer wgD.Done()

			result := Part1(lines, 0, col, 'D')

			chD <- result
		}(i)
	}
	wgD.Wait()
	close(chD)

	for i := range chD {
		if i > max {
			max = i
		}
	}

	var wgU sync.WaitGroup
	chU := make(chan int, width)

	for i := 0; i < width; i++ {
		wgU.Add(1)
		go func(col int) {
			defer wgU.Done()

			result := Part1(lines, height-1, col, 'U')

			chU <- result
		}(i)
	}
	wgU.Wait()
	close(chU)

	for i := range chU {
		if i > max {
			max = i
		}
	}

	var wgL sync.WaitGroup
	chL := make(chan int, height)

	for i := 0; i < height; i++ {
		wgL.Add(1)
		go func(row int) {
			defer wgL.Done()

			result := Part1(lines, row, width-1, 'L')

			chL <- result
		}(i)
	}
	wgL.Wait()
	close(chL)

	for i := range chL {
		if i > max {
			max = i
		}
	}

	var wgR sync.WaitGroup
	chR := make(chan int, height)

	for i := 0; i < height; i++ {
		wgR.Add(1)
		go func(row int) {
			defer wgR.Done()

			result := Part1(lines, row, width-1, 'L')

			chR <- result
		}(i)
	}
	wgR.Wait()
	close(chR)

	for i := range chR {
		if i > max {
			max = i
		}
	}

	return max
}