package main

import (
	"fmt"
	"os"
	"regexp"
	"strings"

	"git.annabunch.es/annabunches/adventofcode/2020/lib/util"
)

const (
	STATE_RULES = iota
	STATE_DATA
)

type Node struct {
	name     string
	children [][]*Node
}

func NewRule(name string) *Node {
	return &Node{
		name:     name,
		children: make([][]*Node, 0),
	}
}

func parseRule(name string, rules map[string]string) *Node {
	ruleString := rules[name]
	rule := NewRule(name)

	// and split into subrules
	for _, text := range strings.Split(ruleString, " | ") {
		// quoted strings are our terminals
		if strings.HasPrefix(text, "\"") {
			rule.name = strings.Trim(text, "\"") // change the rule to its "real" name
			continue
		}
		// everything else are rule indexes, which we turn into child nodes
		childSet := make([]*Node, 0)
		for _, childName := range strings.Split(text, " ") {
			childSet = append(childSet, parseRule(childName, rules))
		}
		rule.children = append(rule.children, childSet)
	}

	return rule
}

func parseInput(input []string) (map[string]string, []string) {
	stateRe := regexp.MustCompile("^([0-9]+): (.*)$")
	state := STATE_RULES
	data := make([]string, 0)

	rules := make(map[string]string) // unparsed rules

	for _, line := range input {
		switch state {
		case STATE_RULES:
			if line == "" {
				state = STATE_DATA
				continue
			}

			reData := stateRe.FindAllStringSubmatch(line, 16)

			// rule name
			rules[reData[0][1]] = reData[0][2]
		case STATE_DATA:
			data = append(data, line)
		}
	}

	return rules, data
}

func createLanguage(rules map[string]string, rootName string) map[string]bool {
	root := parseRule(rootName, rules)

	// now we expand the grammar - generate all possible strings in the language
	rawLanguage := expand(root)
	language := make(map[string]bool)
	for _, term := range rawLanguage {
		language[term] = true
	}

	return language
}

// Expand the list of all possible substrings starting with node
func expand(node *Node) []string {
	items := make([]string, 0)

	if len(node.children) == 0 {
		items = append(items, node.name)
		return items
	}

	for _, childSet := range node.children {
		newItems := make([][]string, 0)
		for _, child := range childSet {
			newItems = append(newItems, expand(child))
		}
		items = append(items, combine(newItems, "")...)
	}

	return items
}

// takes a list of strings and combines them in order, returning a list of
// strings of all possible combinations. Example:
// [["foo"] ["bar baz"] ["fnord"]] => ["foobarfnord" "foobazfnord"]
func combine(terms [][]string, acc string) []string {
	remaining := len(terms)
	results := make([]string, 0)

	if remaining == 0 {
		return results
	}

	for _, part := range terms[0] {
		newAcc := acc + part
		if remaining == 1 {
			results = append(results, newAcc)
			continue
		}

		results = append(results, combine(terms[1:], newAcc)...)
	}

	return results
}

func main() {
	step := os.Args[1]
	values := util.InputParserStrings(os.Args[2])
	rules, data := parseInput(values)
	count := 0

	switch step {
	case "1":
		language := createLanguage(rules, "0")
		for _, item := range data {
			if language[item] {
				count++
			}
		}
	case "2":
		// stub
		// for step 2, our base rule expands to 0: 42{N} 31{M}, for N > M and N > 2
		// we exploit this, along with the fact that all valid strings for 42 and 31 are
		// 8 characters long, to "cheat" a little
		matchLength := 8
		left := createLanguage(rules, "42")
		right := createLanguage(rules, "31")

		for _, item := range data {
			leftCount := 0
			rightCount := 0
			onLeft := true
			valid := true // valid until proven otherwise

			for index := 0; index < len(item); index += matchLength {
				if len(item[index:]) < matchLength {
					fmt.Println("Bad Length")
					valid = false // wrong length alignment
					break
				}
				subString := item[index : index+matchLength]
				if onLeft {
					if left[subString] {
						leftCount++
					} else {
						onLeft = false
					}
				}
				if !onLeft {
					if right[subString] {
						rightCount++
					} else {
						valid = false
						break
					}
				}
			}

			if valid && leftCount > 1 && rightCount > 0 && leftCount > rightCount {
				count++
			}
		}
	}
	fmt.Println(count)
}