package main

import (
	"aoc/utils"
	"fmt"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

type operation int

func (o operation) compute(values []int) int {
	switch o {
	case add:
		sum := 0
		for _, val := range values {
			sum += val
		}
		return sum
	case mul:
		product := 1
		for _, val := range values {
			product *= val
		}
		return product
	}

	return 0
}

const (
	add operation = iota
	mul
)

var allOperationSymbols = []string{"+", "*"}

func operationFromSymbol(sym string) operation {
	switch sym {
	case "+":
		return add
	case "*":
		return mul
	default:
		panic(fmt.Sprintf("wtf is a %s?", sym))
	}
}

type problems struct {
	values     [][]int
	operations []operation
}

func (p *problems) feed(column string) {
	last := string(column[len(column)-1])
	done := false

	if slices.Contains(allOperationSymbols, last) {
		p.operations = append(p.operations, operationFromSymbol(last))
		column = column[:len(column)-1]
		done = true
	}

	val, _ := strconv.Atoi(strings.TrimSpace(column))
	idx := len(p.values) - 1
	p.values[idx] = append(p.values[idx], val)

	if done {
		p.values = append(p.values, []int{})
	}
}

func (p *problems) addLine(line string) (done bool) {
	parts := strings.Split(line, " ")
	parts = slices.DeleteFunc(parts, func(elem string) bool {
		return elem == ""
	})

	if slices.Contains(allOperationSymbols, parts[0]) {
		p.operations = make([]operation, len(parts))
		for idx, sym := range parts {
			p.operations[idx] = operationFromSymbol(sym)
		}
		done = true
	} else {
		if len(p.values) == 0 {
			lenparts := len(parts)
			p.values = make([][]int, lenparts)
			for idx := range lenparts {
				p.values[idx] = []int{}
			}
		}

		for idx, part := range parts {
			num, _ := strconv.Atoi(part)
			p.values[idx] = append(p.values[idx], num)
		}
		done = false
	}

	return done
}

func (p problems) solve() []int {
	solutions := make([]int, len(p.values))

	for idx, values := range p.values {
		op := p.operations[idx]
		solutions[idx] = op.compute(values)
	}

	return solutions
}

func stepOne(input chan string) (int, error) {
	modernProblems := problems{}
	for line := range input {
		done := modernProblems.addLine(line)
		if done {
			break
		}
	}

	modernSolutions := modernProblems.solve()
	sum := 0
	for _, solution := range modernSolutions {
		sum += solution
	}

	return sum, nil
}

func transposeInputChan(input chan string) []string {
	lines := [][]rune{}
	for line := range input {
		lines = append(lines, []rune(line))
	}

	linecount := len(lines)
	columncount := len(lines[0])
	transposed := make([][]rune, columncount)

	for idx := range transposed {
		transposed[idx] = make([]rune, linecount)
	}

	for row, line := range lines {
		for col, char := range line {
			transposed[col][row] = char
		}
	}

	columns := make([]string, len(transposed))
	for idx, col := range transposed {
		columns[idx] = string(col)
	}

	return columns
}

func stepTwo(input chan string) (int, error) {
	transposedInput := transposeInputChan(input)
	slices.Reverse(transposedInput)

	// problem-set with one empty problem.
	modernProblems := problems{
		values: [][]int{[]int{}},
	}

	for _, column := range transposedInput {
		if matched, _ := regexp.MatchString("^\\s*$", column); matched {
			continue
		}

		modernProblems.feed(column)
	}

	// Remove last useless empty problem.
	modernProblems.values = modernProblems.values[:len(modernProblems.values)-1]

	modernSolutions := modernProblems.solve()
	sum := 0
	for _, solution := range modernSolutions {
		sum += solution
	}

	return sum, nil
}

func main() {
	inputFile := utils.FilePath("day06.txt")
	utils.RunStep(utils.ONE, inputFile, stepOne)
	utils.RunStep(utils.TWO, inputFile, stepTwo)
}