package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"

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

type state struct {
	Pos   helpers.Coordinate3d
	Steps int
}

func main() {
	file := "input"
	if helpers.GetArgNumber(1) != "" {
		file = helpers.GetArgNumber(1)
	}
	f, err := os.Open(file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	part := helpers.GetArgNumber(2)
	if part != "2" {
		part1(f)
	} else {
		part2(f)
	}
}

func part1(r io.Reader) {
	vault, doors, keys, start, allkeys := readVault(r)
	fmt.Println(search(vault, doors, keys, start, allkeys, 0))
}

func part2(r io.Reader) {
	vault, doors, keys, start, allKeys := readVault(r)
	directions := []helpers.Coordinate{{X: 0, Y: -1}, {X: 1, Y: 0}, {X: 0, Y: 1}, {X: -1, Y: 0}}
	vault[start] = false
	for _, d := range directions {
		vault[helpers.Coordinate{X: start.X + d.X, Y: start.Y + d.Y}] = false
	}
	total := 0

	haveKeys := allKeys
	for x := 0; x < start.X; x++ {
		for y := 0; y < start.Y; y++ {
			haveKeys ^= keys[helpers.Coordinate{X: x, Y: y}]
		}
	}

	total += search(vault, doors, keys, helpers.Coordinate{X: start.X - 1, Y: start.Y - 1}, allKeys, haveKeys)

	haveKeys = allKeys
	for x := start.X + 1; x <= start.X*2; x++ {
		for y := 0; y < start.Y; y++ {
			haveKeys ^= keys[helpers.Coordinate{X: x, Y: y}]
		}
	}

	total += search(vault, doors, keys, helpers.Coordinate{X: start.X + 1, Y: start.Y - 1}, allKeys, haveKeys)

	haveKeys = allKeys
	for x := start.X + 1; x <= start.X*2; x++ {
		for y := start.Y + 1; y <= start.Y*2; y++ {
			haveKeys ^= keys[helpers.Coordinate{X: x, Y: y}]
		}
	}

	total += search(vault, doors, keys, helpers.Coordinate{X: start.X + 1, Y: start.Y + 1}, allKeys, haveKeys)

	haveKeys = allKeys
	for x := 0; x < start.X; x++ {
		for y := start.Y + 1; y <= start.Y*2; y++ {
			haveKeys ^= keys[helpers.Coordinate{X: x, Y: y}]
		}
	}

	total += search(vault, doors, keys, helpers.Coordinate{X: start.X - 1, Y: start.Y + 1}, allKeys, haveKeys)

	fmt.Println(total)
}

func readVault(r io.Reader) (map[helpers.Coordinate]bool, map[helpers.Coordinate]int, map[helpers.Coordinate]int, helpers.Coordinate, int) {
	scan := bufio.NewScanner(r)
	var start helpers.Coordinate
	var x, y, allkeys int
	v := make(map[helpers.Coordinate]bool)
	doors, keys := make(map[helpers.Coordinate]int), make(map[helpers.Coordinate]int)
	for scan.Scan() {
		x = 0
		for _, c := range scan.Text() {
			if c != '#' {
				v[helpers.Coordinate{X: x, Y: y}] = true
				if c == '@' {
					start = helpers.Coordinate{X: x, Y: y}
				} else if c != '.' {
					if c < 'a' {
						k := 1 << (c - 'A')
						doors[helpers.Coordinate{X: x, Y: y}] = k
						allkeys |= k
					} else {
						k := 1 << (c - 'a')
						keys[helpers.Coordinate{X: x, Y: y}] = k
						allkeys |= k
					}
				}
			}
			x++
		}
		y++
	}
	if err := scan.Err(); err != nil {
		log.Fatal(err)
	}
	return v, doors, keys, start, allkeys
}

func search(vault map[helpers.Coordinate]bool, doors, keys map[helpers.Coordinate]int, start helpers.Coordinate, allKeys, haveKeys int) int {
	directions := []helpers.Coordinate{{X: 0, Y: -1}, {X: 1, Y: 0}, {X: 0, Y: 1}, {X: -1, Y: 0}}
	queue, visited := []state{{Pos: helpers.Coordinate3d{X: start.X, Y: start.Y, Z: haveKeys}}}, make(map[helpers.Coordinate3d]bool)
	var st state
	for {
		st, queue = queue[0], queue[1:]
		if st.Pos.Z&allKeys == allKeys {
			return st.Steps
		}
		visited[st.Pos] = true
		for _, d := range directions {
			next := helpers.Coordinate3d{X: st.Pos.X + d.X, Y: st.Pos.Y + d.Y, Z: st.Pos.Z}
			if !vault[helpers.Coordinate{X: next.X, Y: next.Y}] || visited[next] {
				continue
			}
			door, ok := doors[helpers.Coordinate{X: next.X, Y: next.Y}]
			if ok && next.Z&door != door {
				continue
			}

			key, ok := keys[helpers.Coordinate{X: next.X, Y: next.Y}]
			if ok {
				next.Z |= key
			}

			queue = append(queue, state{Pos: next, Steps: st.Steps + 1})
		}
	}
}