package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

func main() {
	fmt.Println("Part 1")
	fmt.Println(part1(readData("example1.txt")), 140)
	fmt.Println(part1(readData("example2.txt")), 772)
	fmt.Println(part1(readData("example3.txt")), 1930)
	fmt.Println(part1(readData("data.txt")))
	fmt.Println()
	fmt.Println("Part 2")
	fmt.Println(part2(readData("example1.txt")), 80)
	fmt.Println(part2(readData("example2.txt")), 436)
	fmt.Println(part2(readData("example3.txt")), 1206)
	fmt.Println(part2(readData("example4.txt")), 236)
	fmt.Println(part2(readData("example5.txt")), 368)
	fmt.Println(part2(readData("data.txt")))
}

func part2(plots [][]byte) (totalPrice int) {
	visited := make([][]int, len(plots))
	for i, line := range plots {
		visited[i] = make([]int, len(line))
	}

	nextPlotId := 1
	areas := make([]int, 1)
	for y := range len(plots) {
		for x := range len(plots[y]) {
			if visited[y][x] > 0 {
				continue
			}
			area := visitPlot2(plots, visited, nextPlotId, x, y)
			areas = append(areas, area)
			nextPlotId += 1
		}
	}

	sides := make([]int, len(areas))

	for y := range len(plots) + 1 {
		p0, p1 := 0, 0
		for x := range len(plots[0]) + 1 {
			p0Next := getPlotIdOrZero(visited, x, y-1)
			p1Next := getPlotIdOrZero(visited, x, y)
			scanSides(visited, sides, p0, p1, p0Next, p1Next)
			p0, p1 = p0Next, p1Next
		}
	}

	for x := range len(plots[0]) + 1 {
		p0, p1 := 0, 0
		for y := range len(plots) + 1 {
			p0Next := getPlotIdOrZero(visited, x-1, y)
			p1Next := getPlotIdOrZero(visited, x, y)
			scanSides(visited, sides, p0, p1, p0Next, p1Next)
			p0, p1 = p0Next, p1Next
		}
	}

	for i := range areas {
		totalPrice += areas[i] * sides[i]
	}

	return
}

func scanSides(visited [][]int, sides []int, p0, p1, p0Next, p1Next int) {
	if p0Next != p1Next {
		if p0 == p1 {
			sides[p0Next] += 1
			sides[p1Next] += 1
		} else {
			if p0 != p0Next {
				sides[p0Next] += 1
			}
			if p1 != p1Next {
				sides[p1Next] += 1
			}
		}
	}
}

var DIR = [4][2]int{{0, -1}, {-1, 0}, {0, 1}, {1, 0}}

func visitPlot2(plots [][]byte, visited [][]int, plotId int, x, y int) (area int) {
	if visited[y][x] > 0 {
		return
	}

	visited[y][x] = plotId
	area += 1

	p := plots[y][x]

	for _, dir := range DIR {
		other := getPlotOrZero(plots, x+dir[0], y+dir[1])
		if other == p {
			area += visitPlot2(plots, visited, plotId, x+dir[0], y+dir[1])
		}
	}

	return
}

func getPlotOrZero(plots [][]byte, x, y int) byte {
	if x < 0 || y < 0 || y >= len(plots) || x >= len(plots[y]) {
		return 0
	} else {
		return plots[y][x]
	}
}

func getPlotIdOrZero(plots [][]int, x, y int) int {
	if x < 0 || y < 0 || y >= len(plots) || x >= len(plots[y]) {
		return 0
	} else {
		return plots[y][x]
	}
}

func part1(plots [][]byte) (totalPrice int) {
	visited := make([][]bool, len(plots))
	for i, line := range plots {
		visited[i] = make([]bool, len(line))
	}

	for y := range len(plots) {
		for x := range len(plots[y]) {
			area, perimeter := visitPlot1(plots, visited, x, y)
			totalPrice += area * perimeter
		}
	}

	return
}

func visitPlot1(plots [][]byte, visited [][]bool, x, y int) (area, perimeter int) {
	if visited[y][x] {
		return
	}

	visited[y][x] = true
	area += 1

	p := plots[y][x]

	for _, dir := range DIR {
		other := getPlotOrZero(plots, x+dir[0], y+dir[1])

		if other == p {
			otherArea, otherPerimeter := visitPlot1(plots, visited, x+dir[0], y+dir[1])
			area += otherArea
			perimeter += otherPerimeter
		} else {
			perimeter += 1
		}
	}

	return
}

func readData(fileName string) (plots [][]byte) {
	fp, err := os.Open(fileName)
	if err != nil {
		panic(err)
	}

	scanner := bufio.NewScanner(fp)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		plots = append(plots, []byte(line))
	}
	return
}