package main

import (
	"fmt"
	"strings"

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

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

func part1(input []string) {
	var records []string
	var groups [][]int
	for _, ln := range input {
		pts := strings.Fields(ln)
		records = append(records, pts[0])
		var grp []int
		for _, n := range strings.Split(pts[1], ",") {
			grp = append(grp, h.Atoi(n))
		}
		groups = append(groups, grp)
	}
	var result int
	for i := range records {
		// Initialize cache for this row
		cache := make(map[string]int)

		// And solve the row
		result += solveRow(0, 0, records[i], groups[i], cache)
	}

	fmt.Println("# Part 1")
	fmt.Println(result)
}

func part2(input []string) {
	var records []string
	var groups [][]int
	for _, ln := range input {
		pts := strings.Fields(ln)
		// Unfold the paper
		pts[0] = fmt.Sprintf("%s?%s?%s?%s?%s", pts[0], pts[0], pts[0], pts[0], pts[0])
		pts[1] = fmt.Sprintf("%s,%s,%s,%s,%s", pts[1], pts[1], pts[1], pts[1], pts[1])
		records = append(records, pts[0])
		var grp []int
		for _, n := range strings.Split(pts[1], ",") {
			grp = append(grp, h.Atoi(n))
		}
		groups = append(groups, grp)
	}

	var result int
	for i := range records {
		// Initialize cache for this row
		cache := make(map[string]int)

		// And solve the row
		result += solveRow(0, 0, records[i], groups[i], cache)
	}

	fmt.Println("# Part 2")
	fmt.Println(result)
}

func ck(i, j int) string { return fmt.Sprintf("%d_%d", i, j) }

func solveRow(rbIdx, groupIdx int, record string, group []int, cache map[string]int) int {
	if rbIdx >= len(record) {
		if groupIdx < len(group) {
			return 0
		}
		return 1
	}
	if v, ok := cache[ck(rbIdx, groupIdx)]; ok {
		return v
	}
	res := 0
	if record[rbIdx] == '.' {
		res = solveRow(rbIdx+1, groupIdx, record, group, cache)
	} else {
		if record[rbIdx] == '?' {
			res += solveRow(rbIdx+1, groupIdx, record, group, cache)
		}
		if groupIdx < len(group) {
			count := 0
			for k := rbIdx; k < len(record); k++ {
				if count > group[groupIdx] || record[k] == '.' || count == group[groupIdx] && record[k] == '?' {
					break
				}
				count += 1
			}
			if count == group[groupIdx] {
				if rbIdx+count < len(record) && record[rbIdx+count] != '#' {
					res += solveRow(rbIdx+count+1, groupIdx+1, record, group, cache)
				} else {
					res += solveRow(rbIdx+count, groupIdx+1, record, group, cache)
				}
			}
		}
	}
	cache[ck(rbIdx, groupIdx)] = res
	return res
}