package main

import (
	"log"
	"os"
	"strings"

	"github.com/thoas/go-funk"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

func puzzle1(input []string) int {
	return do(input, false)
}

func puzzle2(input []string) int {
	return do(input, true)
}

func do(input []string, puzzle2 bool) int {
	connections := map[string][]string{}
	// Read input
	for _, line := range input {
		split := strings.Split(line, "-")
		connections[split[0]] = append(connections[split[0]], split[1])
		connections[split[1]] = append(connections[split[1]], split[0])
	}
	// Find all routes from start to end
	return len(route(puzzle2, connections, []string{"start"}, false))
}

func route(puzzle2 bool, connections map[string][]string, visited []string, smallTwice bool) (routes [][]string) {
	// Find all routes from last visited to end
	for _, next := range connections[visited[len(visited)-1]] {
		// Check if next is end
		if next == "end" {
			routes = append(routes, append(visited, next))
			continue
		}
		// Can visit a single small cave twice, but all remaining caves once
		if next[0] >= 'a' && next[0] <= 'z' && funk.ContainsString(visited, next) {
			if next != "start" && puzzle2 && !smallTwice {
				routes = append(routes, route(puzzle2, connections, append(visited, next), true)...)
			}
			continue
		}
		// Check next possible routes
		routes = append(routes, route(puzzle2, connections, append(visited, next), smallTwice)...)
	}
	// Return all found routes
	return routes
}