package main

import (
	"AOC2022/helper"
	"fmt"
	"strconv"
	"strings"
)

type State struct {
	currentRessources [4]int
	currentProduction [4]int
	runtime           int
}

type Blueprint struct {
	oreRoboterCost        int
	clayRoboterCost       int
	obsidianRobototerCost [2]int
	geodeRoboterCost      [2]int
}

func main() {
	//args := os.Args[1:]
	lines := helper.ReadTextFile("day19/input")
	highestGeode := make([]int, len(lines))
	for i, line := range lines {
		blueprint := getBluePrint(line)
		fastestTime := getFastestTimeToElementN(blueprint, map[[9]int]State{[9]int{0, 0, 0, 0, 1, 0, 0, 0, 0}: State{[4]int{0, 0, 0, 0}, [4]int{1, 0, 0, 0}, 0}}, 2)
		startStates := getAllPossibleCombinationsWithFastestTime(blueprint, map[[9]int]State{[9]int{0, 0, 0, 0, 1, 0, 0, 0, 0}: State{[4]int{0, 0, 0, 0}, [4]int{1, 0, 0, 0}, 0}}, 2, fastestTime)
		tmpStates := getAllPossibleCombinationsWithFastestTime(blueprint, map[[9]int]State{[9]int{0, 0, 0, 0, 1, 0, 0, 0, 0}: State{[4]int{0, 0, 0, 0}, [4]int{1, 0, 0, 0}, 0}}, 2, fastestTime)
		fastestTime = getFastestTimeToElementN(blueprint, tmpStates, 3)
		startStates = getAllPossibleCombinationsWithFastestTime(blueprint, startStates, 3, fastestTime)
		highestGeode[i] = getHighestGeode(blueprint, startStates)
		fmt.Println(i)
	}
	sum := 0
	for i, score := range highestGeode {
		sum += (i + 1) * score
	}
	fmt.Println(sum)

}

func getHighestGeode(blueprint Blueprint, startStates map[[9]int]State) int {
	activeStates := startStates
	endStates := []State{}
	for len(activeStates) > 0 {
		stepHighestGeode(&activeStates, &endStates, &blueprint)
	}
	highestgeode := 0
	for _, state := range endStates {
		if state.currentRessources[3] > highestgeode && state.runtime >= 25 {
			highestgeode = state.currentRessources[3]
		}
	}
	return highestgeode
}

func getFastestTimeToElementN(blueprint Blueprint, startStates map[[9]int]State, untilElementN int) int {
	activeStates := startStates
	fastestTimeToObsidianState := State{[4]int{0, 0, 0, 0}, [4]int{1, 0, 0, 0}, 26}
	for len(activeStates) > 0 {
		step(&activeStates, &fastestTimeToObsidianState, untilElementN, &blueprint)
	}
	return fastestTimeToObsidianState.runtime
}

func getAllPossibleCombinationsWithFastestTime(blueprint Blueprint, startStates map[[9]int]State, untilElementN int, fastestTime int) map[[9]int]State {
	activeStates := startStates
	fastestTimeToObsidianState := map[[9]int]State{}
	for len(activeStates) > 0 {
		stepFindAllFastestTime(&activeStates, &fastestTimeToObsidianState, untilElementN, fastestTime, &blueprint)
	}
	return fastestTimeToObsidianState
}

func step(activeStates *map[[9]int]State, fastestTImeToObsidian *State, untilElementN int, blueprint *Blueprint) {
	newTmpStates := generatePossibleTmpStates(activeStates, blueprint)
	for _, tmpState := range newTmpStates {
		if tmpState.currentProduction[untilElementN] > 0 && (*fastestTImeToObsidian).runtime > tmpState.runtime {
			*fastestTImeToObsidian = tmpState
		}
		if tmpState.currentProduction[untilElementN] == 0 && tmpState.runtime < (*fastestTImeToObsidian).runtime {
			(*activeStates)[getIdentifier(tmpState)] = tmpState
		}
	}
}

func stepHighestGeode(activeStates *map[[9]int]State, endStates *[]State, blueprint *Blueprint) {
	newTmpStates := generatePossibleTmpStates(activeStates, blueprint)
	for _, tmpState := range newTmpStates {
		if tmpState.runtime < 24 {
			(*activeStates)[getIdentifier(tmpState)] = tmpState
		} else if tmpState.runtime > 24 {
			fmt.Println("WTF")
		} else {
			*endStates = append(*endStates, tmpState)
		}
	}
}

func stepFindAllFastestTime(activeStates *map[[9]int]State, fastestTImeStates *map[[9]int]State, untilElementN int, fastestTime int, blueprint *Blueprint) {
	newTmpStates := generatePossibleTmpStates(activeStates, blueprint)
	for _, tmpState := range newTmpStates {
		if tmpState.currentProduction[untilElementN] > 0 && fastestTime >= tmpState.runtime && tmpState.runtime < 24 {
			identifier := getIdentifier(tmpState)
			(*fastestTImeStates)[identifier] = tmpState
		}
		if tmpState.currentProduction[untilElementN] == 0 && tmpState.runtime < fastestTime && tmpState.runtime < 24 {
			(*activeStates)[getIdentifier(tmpState)] = tmpState
		}
	}
}

func getIdentifier(tmpState State) [9]int {
	identifier := [9]int{}
	copy(identifier[:], append(append(tmpState.currentRessources[:], tmpState.currentProduction[:]...), tmpState.runtime)[:9])
	return identifier
}

func generatePossibleTmpStates(activeStates *map[[9]int]State, blueprint *Blueprint) map[[9]int]State {
	key := get_some_key(*activeStates)
	activeState := (*activeStates)[key]
	delete(*activeStates, key)
	possibleProductions := activeState.getPossibleProductions(blueprint)
	newTmpStates := make(map[[9]int]State)
	for i := -1; i < len(possibleProductions); i++ {
		if i == -1 || possibleProductions[i] == 1 {
			tmpState := activeState
			tmpState.produceRessources()
			tmpState.produceRoboter(blueprint, i)
			tmpState.runtime++
			newTmpStates[getIdentifier(tmpState)] = tmpState
		}
	}
	return newTmpStates
}

func (state *State) produceRoboter(blueprint *Blueprint, roboter int) {
	switch roboter {
	case 0:
		state.currentProduction[0]++
		state.currentRessources[0] -= blueprint.oreRoboterCost
	case 1:
		state.currentProduction[1]++
		state.currentRessources[0] -= blueprint.clayRoboterCost
	case 2:
		state.currentProduction[2]++
		state.currentRessources[0] -= blueprint.obsidianRobototerCost[0]
		state.currentRessources[1] -= blueprint.obsidianRobototerCost[1]
	case 3:
		state.currentProduction[3]++
		state.currentRessources[0] -= blueprint.geodeRoboterCost[0]
		state.currentRessources[2] -= blueprint.geodeRoboterCost[1]
	default:

	}
}

func (state *State) produceRessources() {
	for i := 0; i < 4; i++ {
		state.currentRessources[i] += state.currentProduction[i]
	}
}

func (state State) getPossibleProductions(blueprint *Blueprint) [4]int {
	possibleProductions := [4]int{0, 0, 0, 0}
	if state.currentRessources[0] >= blueprint.oreRoboterCost {
		possibleProductions[0] = 1
	}
	if state.currentRessources[0] >= blueprint.clayRoboterCost {
		possibleProductions[1] = 1
	}
	if state.currentRessources[0] >= blueprint.obsidianRobototerCost[0] && state.currentRessources[1] >= blueprint.obsidianRobototerCost[1] {
		possibleProductions[2] = 1
	}
	if state.currentRessources[0] >= blueprint.geodeRoboterCost[0] && state.currentRessources[2] >= blueprint.geodeRoboterCost[1] {
		possibleProductions[3] = 1
	}
	return possibleProductions
}

func getBluePrint(line string) Blueprint {
	productionCostStrings := strings.Split(strings.Split(line, ":")[1], ". ")
	oreCostString := strings.ReplaceAll(productionCostStrings[0], " ore", "")[18:]
	oreCost, _ := strconv.Atoi(oreCostString)
	clayCostString := strings.ReplaceAll(productionCostStrings[1], " ore", "")[22:]
	claycost, _ := strconv.Atoi(clayCostString)
	obsidiancostString := strings.ReplaceAll(productionCostStrings[2], "ore and ", "")[26:]
	obsidianCost := helper.StringSliceToIntSlice(strings.Split(obsidiancostString[:len(obsidiancostString)-5], " "))
	geodeCostString := strings.ReplaceAll(productionCostStrings[3], "ore and ", "")[23:]
	geodeCost := helper.StringSliceToIntSlice(strings.Split(geodeCostString[:len(geodeCostString)-10], " "))
	return Blueprint{oreCost, claycost, [2]int{obsidianCost[0], obsidianCost[1]}, [2]int{geodeCost[0], geodeCost[1]}}
}

func get_some_key(m map[[9]int]State) [9]int {
	for k := range m {
		return k
	}
	return [9]int{}
}