package main

import (
	"log"
	"math"
	"os"
	"strconv"
	"strings"

	"github.com/thoas/go-funk"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := string(inputBytes)

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

func puzzle1(input string) int {
	return do(input, false)
}

func puzzle2(input string) int {
	return do(input, true)
}

func do(input string, puzzle2 bool) int {
	// Read input
	crabs := []int{}
	for _, p := range strings.Split(input, ",") {
		i, _ := strconv.Atoi(p)
		crabs = append(crabs, i)
	}
	// Get min and max
	minPos := funk.MinInt(crabs)
	maxPos := funk.MaxInt(crabs)
	// For each possible target position, calculate required fuel
	minFuel := math.MaxInt
	for target := minPos; target <= maxPos; target++ {
		fuel := 0
		for _, c := range crabs {
			if puzzle2 {
				if c == target {
					continue
				}
				newFuel := 0
				distance := funk.ShortIf(c > target, c-target, target-c).(int)
				for i := 1; i <= distance; i++ {
					// First step costs 1 fuel, second step costs 2 fuel, etc.
					newFuel += i
				}
				fuel += newFuel
			} else {
				// Add distance to fuel
				fuel += int(math.Abs(float64(c) - float64(target)))
			}
		}
		// Check if we have a new minimum
		if fuel < minFuel {
			minFuel = fuel
		}
	}
	return minFuel
}