package day09

func Part1(input string) int {
	files, gaps, lengthOfFinalFile := parseInput(input)
	checksum := compactFiles(files, gaps, lengthOfFinalFile)
	return checksum
}

func Part2(input string) int {
	files, freeSpaces := parseInputPart2(input)
	checksum := calculateChecksum(files, freeSpaces)
	return checksum
}

func parseInput(input string) ([]int, []int, int) {
	var files []int
	var gaps []int
	lengthOfFinalFile := 0
	for i, char := range input {
		if i%2 == 0 {
			files = append(files, int(char-'0'))
			lengthOfFinalFile += int(char - '0')
		} else {
			gaps = append(gaps, int(char-'0'))
		}
	}
	return files, gaps, lengthOfFinalFile
}

func parseInputPart2(input string) ([][]int, [][]int) {
	files, freeSpaces, position := [][]int{}, [][]int{}, 0
	for index, char := range input {
		length := int(char - '0')
		if index%2 == 0 {
			files = append(files, generateRange(position, position+length))
		} else {
			freeSpaces = append(freeSpaces, generateRange(position, position+length))
		}
		position += length
	}
	return files, freeSpaces
}

func compactFiles(files []int, gaps []int, lengthOfFinalFile int) int {
	fileIndex := len(files) - 1
	gapIndex := 0
	position := files[0]
	checksum := 0
	fileID := files[0]

	for lengthOfFinalFile != position {
		// If there's a gap and a file to move then move the file
		if gaps[gapIndex] > 0 && files[fileIndex] > 0 {
			files[fileIndex]--
			gaps[gapIndex]--
			checksum += fileIndex * position
			position++
		}

		// If this file is fully moved go to next file
		if files[fileIndex] == 0 {
			fileIndex--
			fileID++
		}

		// Update gapIndex if this gap is filled
		if gaps[gapIndex] == 0 {
			gapIndex++
			for f := 0; f < files[gapIndex]; f++ {
				checksum += gapIndex * position
				position++
			}
		}
	}

	return checksum
}

func generateRange(start, end int) []int {
	rangeList := make([]int, end-start)
	for i := range rangeList {
		rangeList[i] = start + i
	}
	return rangeList
}

func calculateChecksum(files, freeSpaces [][]int) int {
	checksum := 0
	for fileIndex := len(files) - 1; fileIndex >= 0; fileIndex-- {
		for spaceIndex := 0; spaceIndex < len(freeSpaces); spaceIndex++ {
			if len(freeSpaces[spaceIndex]) >= len(files[fileIndex]) && files[fileIndex][0] > freeSpaces[spaceIndex][0] {
				files[fileIndex] = freeSpaces[spaceIndex][:len(files[fileIndex])]
				freeSpaces[spaceIndex] = freeSpaces[spaceIndex][len(files[fileIndex]):]
			}
		}
	}

	for fileID, file := range files {
		for _, block := range file {
			checksum += fileID * block
		}
	}
	return checksum
}