package day19

import (
	"strings"
	"tools"
)

func Part1(puzzle tools.AoCPuzzle) interface{} {
	input := puzzle.GetInputArray()
	molecule := input[len(input)-1]
	replacementStrings := input[:len(input)-2]
	outcomes := tools.NewSet()

	for _, line := range replacementStrings {
		parts := strings.Split(line, " => ")
		doubleChar := len(parts[0]) == 2
		var result string
		for i, c := range molecule {
			if doubleChar && i == len(molecule)-1 {
				break
			}
			if doubleChar {
				if byte(c) == parts[0][0] && molecule[i+1] == parts[0][1] {
					result = molecule[:i] + parts[1]
					if i < len(molecule)-2 {
						result += molecule[i+2:]
					}
					outcomes.Add(result)
				}
			} else {
				if string(c) == parts[0] {
					result = molecule[:i] + parts[1]
					if i < len(molecule)-1 {
						result += molecule[i+1:]
					}
					outcomes.Add(result)
				}
			}
		}
	}

	return len(outcomes)
}

func getPossibleTransformations(from tools.Set, replacments map[string]tools.Set, maxLen int) tools.Set {
	returnSet := tools.NewSet()
	var result string

	for interfaceMolecule := range from {
		currentMolecule := interfaceMolecule.(string)
		for i := range currentMolecule {
			if transSet, ok := replacments[string(currentMolecule[i])]; ok {
				for replacement := range transSet {
					result = currentMolecule[:i] + replacement.(string)
					if i < len(currentMolecule)-1 {
						result += currentMolecule[i+1:]
					}
					if len(result) <= maxLen {
						returnSet.Add(result)
					}
				}
			}
			if i < len(currentMolecule)-1 {
				if transSet, ok := replacments[currentMolecule[i:i+2]]; ok {
					for replacment := range transSet {
						result = currentMolecule[:i] + replacment.(string)
						if i < len(currentMolecule)-2 {
							result += currentMolecule[i+2:]
						}
						if len(result) <= maxLen {
							returnSet.Add(result)
						}
					}
				}
			}
		}
	}

	return returnSet
}

func Part2(puzzle tools.AoCPuzzle) interface{} {
	input := puzzle.GetInputArray()
	molecule := input[len(input)-1]
	replacementStrings := input[:len(input)-2]
	replacements := make(map[string]tools.Set)
	for _, line := range replacementStrings {
		parts := strings.Split(line, " => ")
		if _, ok := replacements[parts[0]]; !ok {
			replacements[parts[0]] = tools.NewSet()
		}
		replacements[parts[0]].Add(parts[1])
	}

	currentMolecules := tools.NewSet()
	currentMolecules.Add("e")
	transformCount := 0
	for !currentMolecules.Contains(molecule) {
		transformCount++
		currentMolecules = getPossibleTransformations(currentMolecules, replacements, len(molecule))
	}

	return transformCount
}