package main

import (
	"fmt"
	"log"
	"os"
	"strconv"

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

// operators
const (
	ADD = -1
	MUL = -2
	GRP = -3
)

type Stack struct {
	data []int
}

func NewStack() *Stack {
	return &Stack{
		data: make([]int, 0),
	}
}

func (s *Stack) Push(value int) {
	s.data = append(s.data, value)
}

func (s *Stack) Pop() int {
	if s.Empty() {
		log.Panicf("Tried to pop from empty stack")
	}
	ret := s.Top()
	s.data = s.data[:len(s.data)-1]
	return ret
}

func (s *Stack) Top() int {
	if s.Empty() {
		return -3
	}
	return s.data[len(s.data)-1]
}

func (s *Stack) Empty() bool {
	return len(s.data) == 0
}

// This only works when the starting expression is all positive, single-digit integers
// and where all operators are additive (i.e. we rely on being able to use negative numbers
// as sentinel values.
func convertRPN(exp string, precedence map[int]int) []int {
	rpn := make([]int, 0)
	ops := NewStack()
	for _, i := range exp {
		e := string(i)
		x, err := strconv.Atoi(e)
		if err == nil {
			rpn = append(rpn, x)
			continue
		}

		var opcode int
		switch e {
		case "+":
			opcode = ADD
		case "*":
			opcode = MUL
		case "(":
			ops.Push(GRP)
			continue
		case ")":
			for ops.Top() != GRP {
				rpn = append(rpn, ops.Pop())
			}
			ops.Pop()
			continue
		default:
			continue
		}

		if ops.Empty() {
			ops.Push(opcode)
			continue
		}
		if precedence[ops.Top()] >= precedence[opcode] {
			rpn = append(rpn, ops.Pop())
		}
		ops.Push(opcode)
	}
	for !ops.Empty() {
		rpn = append(rpn, ops.Pop())
	}

	return rpn
}

func evaluateExpression(rpn []int) int {
	output := NewStack()
	for _, value := range rpn {
		if value >= 0 {
			output.Push(value)
			continue
		}
		switch value {
		case ADD:
			x := output.Pop()
			y := output.Pop()
			output.Push(x + y)
		case MUL:
			x := output.Pop()
			y := output.Pop()
			output.Push(x * y)
		}
	}

	return output.Pop()
}

func main() {
	step := os.Args[1]
	values := util.InputParserStrings(os.Args[2])

	total := 0
	precedence := make(map[int]int)
	switch step {
	case "1":
		precedence[ADD] = 0
		precedence[MUL] = 0
		precedence[GRP] = -1
	case "2":
		precedence[ADD] = 1
		precedence[MUL] = 0
		precedence[GRP] = -1
	}

	exps := make([][]int, 0)
	for _, line := range values {
		// parse the expression into RPN
		exps = append(exps, convertRPN(line, precedence))
	}

	for _, exp := range exps {
		total += evaluateExpression(exp)
	}
	fmt.Println(total)
}