package day11

import (
	_ "fmt"
	"strings"
	"sort"
	"adventofcode2022/utils"
	"math/big"
)

type op int

const (
	addition op = iota + 1
	multiply
	square
)

type monkey struct {
	items []*big.Int
	operation op
	opArg *big.Int
	test *big.Int
	ifTrue int
	ifFalse int
	inspected int
}

func Part1(input string) int {

	monkeys := parseInput(input)
	var level = big.NewInt(0)
	for round := 0; round < 20; round++ {
		for i, monkey := range monkeys {
			for _, item := range monkey.items {
				monkeys[i].inspected++
				switch monkey.operation {
				case addition:
					level = new(big.Int).Add(item, monkey.opArg)
				case multiply:
					level = new(big.Int).Mul(item, monkey.opArg)
				case square:
					level = new(big.Int).Mul(item, item)
				}
				level = new(big.Int).Div(level, big.NewInt(3))
				t := &big.Int{}
				t.Mod(level, monkey.test)
				if t.Sign() == 0 {
					monkeys[monkey.ifTrue].items = append(monkeys[monkey.ifTrue].items, level)
				} else {
					monkeys[monkey.ifFalse].items = append(monkeys[monkey.ifFalse].items, level)
				}
				
			}
			monkeys[i].items = monkeys[i].items[:0]
		}
	}
	sort.Slice(monkeys, func(i, j int) bool {
		return monkeys[i].inspected > monkeys[j].inspected
	})
	
	return  monkeys[0].inspected * monkeys[1].inspected
}

func Part2(input string) int {
	monkeys := parseInput(input)
	var level = big.NewInt(0)
	commonDivisor := big.NewInt(3 * 13 * 19 * 17 * 5 * 7 * 11 * 2)
	for round := 0; round < 10000; round++ {
		for i, monkey := range monkeys {
			for _, item := range monkey.items {
				monkeys[i].inspected++
				switch monkey.operation {
				case addition:
					level = new(big.Int).Add(item, monkey.opArg)
				case multiply:
					level = new(big.Int).Mul(item, monkey.opArg)
				case square:
					level = new(big.Int).Mul(item, item)
				}
				level = level.Mod(level, commonDivisor)
				t := &big.Int{}
				t.Mod(level, monkey.test)
				if t.Sign() == 0 {
					monkeys[monkey.ifTrue].items = append(monkeys[monkey.ifTrue].items, level)
				} else {
					monkeys[monkey.ifFalse].items = append(monkeys[monkey.ifFalse].items, level)
				}
				
			}
			monkeys[i].items = monkeys[i].items[:0]
		}
	}
	sort.Slice(monkeys, func(i, j int) bool {
		return monkeys[i].inspected > monkeys[j].inspected
	})
	
	return  monkeys[0].inspected * monkeys[1].inspected
}

func parseInput(input string) []monkey {
	monkeys := []monkey{}

	for i, monkeyString := range strings.Split(input, "\n\n") {
		for _, line := range strings.Split(monkeyString, "\n") {
			if strings.HasPrefix(line, "Monkey") {
				monkeys = append(monkeys, monkey{})
			} else if strings.HasPrefix(line, "  Starting items: ") {
				s := strings.TrimPrefix(line, "  Starting items: ")
				for _, piece := range strings.Split(s, ", ") {
					t := utils.MustAtoi(piece)
					monkeys[i].items =  append(monkeys[i].items, big.NewInt(int64(t)))
				}
			} else if strings.HasPrefix(line, "  Operation: new = ") {
				s := strings.TrimPrefix(line, "  Operation: new = ")
				if strings.HasPrefix(s, "old + ") {
					s2 := strings.TrimPrefix(s, "old + ")
					monkeys[i].operation = addition
					t := utils.MustAtoi(s2)
					monkeys[i].opArg = big.NewInt(int64(t))
				} else if s == "old * old" {
					monkeys[i].operation = square
				} else if strings.HasPrefix(s, "old * ") {
					s2 := strings.TrimPrefix(s, "old * ")
					monkeys[i].operation = multiply
					t := utils.MustAtoi(s2)
					monkeys[i].opArg = big.NewInt(int64(t))
				}

			} else if strings.HasPrefix(line, "  Test: divisible by ") {
				s := strings.TrimPrefix(line, "  Test: divisible by ")
				t := utils.MustAtoi(s)
				monkeys[i].test = big.NewInt(int64(t))
			} else if strings.HasPrefix(line, "    If true: throw to monkey ") {
				s := strings.TrimPrefix(line, "    If true: throw to monkey ")
				monkeys[i].ifTrue = utils.MustAtoi(s)
			} else if strings.HasPrefix(line, "    If false: throw to monkey ") {
				s := strings.TrimPrefix(line, "    If false: throw to monkey ")
				monkeys[i].ifFalse = utils.MustAtoi(s)
			} else if line == "" {
				// do nothing
			} else {
				panic("unknown input")
			}
		}
	}
	return monkeys
}