package main

import (
	"log"
	"os"
	"strconv"
	"strings"
)

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

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

func puzzle1(input []string) int {
	grid := fold(input, false)
	count := 0
	for _, row := range grid {
		for _, cell := range row {
			if cell == 1 {
				count++
			}
		}
	}
	return count
}

func puzzle2(input []string) string {
	grid := fold(input, true)
	res := ""
	for _, row := range grid {
		for _, cell := range row {
			if cell == 1 {
				res += "#"
			} else {
				res += " "
			}
		}
		res += "\n"
	}
	return res
}

func fold(input []string, puzzle2 bool) [][]int {
	// Parse input
	xs := []int{}
	ys := []int{}
	folds := []string{}
	for _, line := range input {
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "fold along") {
			folds = append(folds, strings.TrimPrefix(line, "fold along "))
			continue
		}
		parts := strings.Split(line, ",")
		x, _ := strconv.Atoi(parts[0])
		xs = append(xs, x)
		y, _ := strconv.Atoi(parts[1])
		ys = append(ys, y)
	}
	// Get maximum x and y
	var maxX, maxY int
	for _, x := range xs {
		if x > maxX {
			maxX = x
		}
	}
	for _, y := range ys {
		if y > maxY {
			maxY = y
		}
	}
	// Create grid
	grid := make([][]int, maxY+1)
	for i := range grid {
		grid[i] = make([]int, maxX+1)
	}
	// Fill grid
	for i, x := range xs {
		grid[ys[i]][x] = 1
	}
	// Fold grid
	for fi, f := range folds {
		if fi > 0 && !puzzle2 {
			break
		}
		parts := strings.Split(f, "=")
		p, _ := strconv.Atoi(parts[1])
		xl := len(grid[0])
		yl := len(grid)
		if parts[0] == "x" {
			// Copy all points from right of the fold to the left and shrink the grid
			for j := 0; j < yl; j++ {
				for i := 0; i+p < xl && p-i >= 0; i++ {
					if grid[j][p+i] == 1 {
						grid[j][p-i] = 1
					}
				}
				// Shrink row
				grid[j] = grid[j][:p]
			}
		} else {
			// Copy all points from bottom of the fold to the top
			for j := 0; j+p < yl && p-j >= 0; j++ {
				for i := 0; i < xl; i++ {
					if grid[p+j][i] == 1 {
						grid[p-j][i] = 1
					}
				}
			}
			// Shrink grid
			grid = grid[:p]
		}
	}
	return grid
}