package main

import (
	"fmt"
	"sync"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := h.StdinToStringSlice()
	part1(inp)
	fmt.Println()
	part2(inp)
}

func part1(inp []string) {
	var ret int
	m := h.StringSliceToCoordByteMap(inp)
	ok := true
	for ok {
		var splits int
		splits, ok = progBeams(&m)
		ret += splits
	}
	fmt.Println("# Part 1")
	fmt.Println(ret)
}

func progBeams(m *h.CoordByteMap) (int, bool) {
	var splits int
	beam, err := m.FindLast('|')
	if err != nil {
		beam, _ = m.FindFirst('S')
	}
	if beam.Y == m.BRY {
		return 0, false
	}
	beams := m.FindAllOnRow(beam.Y, '|')
	if len(beams) == 0 {
		beams = append(beams, beam)
	}

	for i := range beams {
		nxt := m.Get(beams[i].South())
		switch nxt {
		case '.':
			m.Put(beams[i].South(), '|')
		case '^':
			splits++
			m.Put(beams[i].South().East(), '|')
			m.Put(beams[i].South().West(), '|')
		}
	}
	return splits, true
}

func part2(inp []string) {
	m := h.StringSliceToCoordByteMap(inp)
	fmt.Println("# Part 2")
	fmt.Println(quantumRevKickoff(&m))
}

// quantumRevKickoff starts checks for each potential beam output
func quantumRevKickoff(m *h.CoordByteMap) int {
	var ret int
	var wg sync.WaitGroup
	status := make([]byte, m.BRX)
	for i := range status {
		status[i] = '.'
	}
	track := &WorldTracker{
		t: make(map[h.Coordinate]int),
	}
	for x := m.TLX; x < m.BRX; x++ {
		c := h.Coordinate{X: x, Y: m.BRY}
		wg.Go(func() {
			if r, v := quantumRevCalc(m, c, track); v {
				ret += r
				status[c.X] = 'O'
			} else {
				status[c.X] = 'X'
			}
		})
	}
	wg.Wait()
	return ret + 1
}

type WorldTracker struct {
	t    map[h.Coordinate]int
	lock sync.Mutex
}

func (t *WorldTracker) Get(k h.Coordinate) (int, bool) {
	t.lock.Lock()
	defer t.lock.Unlock()
	v, ok := t.t[k]
	return v, ok
}
func (t *WorldTracker) Set(k h.Coordinate, v int) {
	t.lock.Lock()
	defer t.lock.Unlock()
	t.t[k] = v
}

// quantumRevCalc calculates how many worlds would provide a beam at
// the given coordinate
func quantumRevCalc(m *h.CoordByteMap, c h.Coordinate, track *WorldTracker) (int, bool) {
	if v, ok := track.Get(c); ok {
		return v, ok
	}
	var ret int
	var valid bool
	if m.Get(c.East()) == '^' {
		if r, v := quantumRevCalc(m, c.East(), track); v {
			ret += r
			valid = true
		}
	}
	if m.Get(c.West()) == '^' {
		if r, v := quantumRevCalc(m, c.West(), track); v {
			ret += r
			valid = true
		}
	}
	if m.Get(c.North()) == '.' {
		if r, v := quantumRevCalc(m, c.North(), track); v {
			ret += r
			valid = true
		}
	}
	track.Set(c, ret)
	if m.Get(c.North()) == 'S' {
		return 1, true
	}
	return ret, valid
}