package main

import (
	"aoc/utils"
	"fmt"
	"strings"
)

func main() {
	g := NewGraph(utils.ReadLines())
	fmt.Println("part1: ", g.Part1())
	fmt.Println("part2: ", g.Part2())
}

type Graph struct {
	edges    [][]uint16
	bigCaves uint16
}

func (g *Graph) Part1() int {
	initCache()
	return g.countPathsPart1(startCave, 0)
}

func (g *Graph) Part2() int {
	initCache()
	return g.countPathsPart2(startCave, 0)
}

const (
	startCave = 1
	endCave   = 0
)

func caveMask(cave uint16) uint16 {
	return 1 << cave
}

func (g *Graph) countPathsPart1(from uint16, visited uint16) int {
	var count int
	if count, ok := getCachedResult(from, visited); ok {
		return count
	}
	for _, to := range g.edges[from] {
		if to == endCave {
			count++
			continue
		}
		toMask := caveMask(to) &^ g.bigCaves
		if visited&toMask == 0 {
			count += g.countPathsPart1(to, visited|toMask)
		}
	}
	putToCache(from, visited, count)
	return count
}

const doubleMask = 1 << 15

func (g *Graph) countPathsPart2(from uint16, visited uint16) int {
	var count int
	if count, ok := getCachedResult(from, visited); ok {
		return count
	}
	for _, to := range g.edges[from] {
		if to == endCave {
			count++
			continue
		}
		toMask := caveMask(to) &^ g.bigCaves
		if visited&toMask == 0 {
			count += g.countPathsPart2(to, visited|toMask)
		} else if visited&doubleMask == 0 {
			count += g.countPathsPart2(to, visited|doubleMask)
		}
	}
	putToCache(from, visited, count)
	return count
}

// Memoïzation

type key struct {
	cave    uint16
	visited uint16
}

type cache map[key]int

var memo cache

func initCache() {
	memo = make(cache)
}

func getCachedResult(from uint16, visited uint16) (int, bool) {
	count, ok := memo[key{from, visited}]
	return count, ok
}

func putToCache(from uint16, visited uint16, count int) {
	memo[key{from, visited}] = count
}

// Input parsing

func NewGraph(input []string) *Graph {
	g := newBuilder()
	g.parse(input)
	return &g.Graph
}

func newBuilder() *graphBuilder {
	return &graphBuilder{
		Graph: Graph{edges: make([][]uint16, 2, 16)},
		caves: map[string]uint16{
			"start": startCave,
			"end":   endCave,
		},
	}
}

type graphBuilder struct {
	Graph
	caves map[string]uint16
}

func (g *graphBuilder) parse(input []string) {
	for _, line := range input {
		edge := strings.Split(line, "-")
		c1 := g.getCave(edge[0])
		c2 := g.getCave(edge[1])
		g.addEdge(c1, c2)
	}
}

func (g *graphBuilder) addEdge(c1, c2 uint16) {
	if c1 != endCave && c2 != startCave {
		g.edges[c1] = append(g.edges[c1], c2)
	}
	if c1 != startCave && c2 != endCave {
		g.edges[c2] = append(g.edges[c2], c1)
	}
}

func (g *graphBuilder) getCave(name string) uint16 {
	if cave, ok := g.caves[name]; ok {
		return cave
	}
	cave := uint16(len(g.caves))
	g.caves[name] = cave
	if 'A' <= name[0] && name[0] < 'Z' {
		g.bigCaves |= caveMask(cave)
	}
	g.edges = append(g.edges, nil)
	return cave
}