package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"iter"
	"log"
	"os"
	"strconv"
	"strings"
)

func main() {
	f := os.Stdin
	if len(os.Args) > 1 && os.Args[1] != "-" {
		var err error
		f, err = os.Open(os.Args[1])
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
	}

	lines, err := readLines(f)
	if err != nil {
		log.Fatal(err)
	}

	problems1, err := parseProblemsPart1(lines)
	if err != nil {
		log.Fatal(err)
	}
	problems2, err := parseProblemsPart2(lines)
	if err != nil {
		log.Fatal(err)
	}

	var result1, result2 int
	for _, p := range problems1 {
		result1 += p.Solve()
	}
	for _, p := range problems2 {
		result2 += p.Solve()
	}
	fmt.Println("part one:", result1)
	fmt.Println("part two:", result2)
}

type Action int

const (
	ActionAdd Action = iota
	ActionMultiply
)

type problem struct {
	Values []int
	Action
}

func (p problem) Solve() int {
	if len(p.Values) == 0 {
		return 0
	}
	answer := p.Values[0]
	for _, v := range p.Values[1:] {
		switch p.Action {
		case ActionAdd:
			answer += v
		case ActionMultiply:
			answer *= v
		default:
			panic(fmt.Sprintf("unsupported action %d", p.Action))
		}
	}
	return answer
}

func readLines(r io.Reader) ([]string, error) {
	var lines []string
	sc := bufio.NewScanner(r)
	for sc.Scan() {
		lines = append(lines, sc.Text())
	}
	return lines, sc.Err()
}

func parseFields(lines []string) [][]string {
	cutAll := func(lines []string, sep string) (before, after []string, ok bool) {
		var found bool
		var i int
		for _, line := range lines {
			j := strings.Index(line, sep)
			if j < 0 {
				j = len(line)
			} else {
				found = true
			}
			i = max(i, j)
		}
		if !found {
			return lines, nil, false
		}

		before = make([]string, 0, len(lines))
		after = make([]string, 0, len(lines))
		for _, line := range lines {
			before = append(before, line[:min(i, len(line))])
			after = append(after, line[min(i+len(sep), len(line)):])
		}
		return before, after, true
	}

	fields := make([][]string, len(lines))
	for {
		before, after, ok := cutAll(lines, " ")
		if !ok {
			break
		}
		for i, v := range before {
			fields[i] = append(fields[i], v)
		}
		lines = after
	}
	return fields
}

func parseProblemsPart1(lines []string) ([]problem, error) {
	if len(lines) == 0 {
		return nil, nil
	}
	var data [][]string
	for _, line := range lines {
		data = append(data, strings.Fields(line))
	}

	problems := make([]problem, len(data[0]))
	for i, row := range data {
		if len(row) != len(problems) {
			return nil, fmt.Errorf("line %d: unexpected number of fields", i+2)
		}

		last := i == len(data)-1
		for j, v := range row {
			p := &problems[j]
			if last {
				switch v {
				case "+":
					p.Action = ActionAdd
				case "*":
					p.Action = ActionMultiply
				default:
					return nil, fmt.Errorf("unexpected action %q", v)
				}
			} else {
				n, err := strconv.Atoi(v)
				if err != nil {
					return nil, err
				}
				p.Values = append(p.Values, n)
			}
		}
	}
	return problems, nil
}

func parseProblemsPart2(lines []string) ([]problem, error) {
	fields := parseFields(lines)
	var problems []problem

	for col := range columnsSeq(fields) {
		var p problem
		switch s := strings.TrimSpace(col[len(col)-1]); s {
		case "+":
			p.Action = ActionAdd
		case "*":
			p.Action = ActionMultiply
		default:
			return nil, fmt.Errorf("unexpected action %q", s)
		}

		var colBytes [][]byte
		for _, v := range col[:len(col)-1] {
			colBytes = append(colBytes, []byte(v))
		}
		for v := range columnsSeq(colBytes) {
			s := string(bytes.TrimSpace(v))
			n, err := strconv.Atoi(s)
			if err != nil {
				return nil, err
			}
			p.Values = append(p.Values, n)
		}

		problems = append(problems, p)
	}
	return problems, nil
}

func columnsSeq[T any](data [][]T) iter.Seq[[]T] {
	return func(yield func([]T) bool) {
		if len(data) < 1 {
			return
		}
		for x := range len(data[0]) {
			var col []T
			for y := range data {
				col = append(col, data[y][x])
			}
			if !yield(col) {
				break
			}
		}
	}
}