package main

import (
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"

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

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

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

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

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

type lineEnd struct {
	x, y int
}

type line struct {
	start, end lineEnd
}

func do(input []string, puzzle2 bool) int {
	// Parse input into lines of form "x1,y1 -> x2,y2"
	lineRegex := regexp.MustCompile(`(\d+),(\d+)\s\->\s(\d+),(\d+)`)
	lines := []line{}
	for _, inputRow := range input {
		lineMatch := lineRegex.FindStringSubmatch(inputRow)
		x1, _ := strconv.Atoi(lineMatch[1])
		y1, _ := strconv.Atoi(lineMatch[2])
		x2, _ := strconv.Atoi(lineMatch[3])
		y2, _ := strconv.Atoi(lineMatch[4])
		lines = append(lines, line{
			start: lineEnd{x1, y1},
			end:   lineEnd{x2, y2},
		})
	}
	// Get maximum x and y values
	maxX, maxY := 0, 0
	for _, l := range lines {
		if l.start.x > maxX {
			maxX = l.start.x
		}
		if l.start.y > maxY {
			maxY = l.start.y
		}
		if l.end.x > maxX {
			maxX = l.end.x
		}
		if l.end.y > maxY {
			maxY = l.end.y
		}
	}
	// Create a 2D array of zeros
	diagram := make([][]int, maxY+1)
	for i := range diagram {
		diagram[i] = make([]int, maxX+1)
	}
	// For each line, increment the value of the cells it crosses
	for _, l := range lines {
		if l.start.x == l.end.x || l.start.y == l.end.y {
			// Vertical line or horizontal line
			for x := funk.MinInt([]int{l.start.x, l.end.x}); x <= funk.MaxInt([]int{l.start.x, l.end.x}); x++ {
				for y := funk.MinInt([]int{l.start.y, l.end.y}); y <= funk.MaxInt([]int{l.start.y, l.end.y}); y++ {
					diagram[y][x]++
				}
			}
		} else if puzzle2 && (l.start.x-l.end.x == l.start.y-l.end.y || l.start.x-l.end.x == l.end.y-l.start.y) {
			// 45-degree diagonal line
			xDir := funk.ShortIf(l.start.x > l.end.x, -1, 1).(int)
			yDir := funk.ShortIf(l.start.y > l.end.y, -1, 1).(int)
			x := l.start.x
			for y := l.start.y; y != l.end.y+yDir; y += yDir {
				diagram[y][x]++
				x += xDir
			}
		}
	}
	// Count the number of cells that have more than one line crossing them
	count := 0
	for _, row := range diagram {
		for _, cell := range row {
			if cell > 1 {
				count++
			}
		}
	}
	return count
}