package aoc2024

import println
import readInput

/*
--- Day 9: Disk Fragmenter ---
https://adventofcode.com/2024/day/9
*/
fun main() {

    val inlineTestInput = """2333133121414131402"""

    fun part1(input: List<String>): Long {
        val diskSize = input[0].sumOf { it - '0' }
        val bitmap = IntArray(diskSize) { -1 }
        var pos = 0
        var isFile = true
        var fileId = 0
        var usedBlocks = 0
        for (c in input[0]) {
            val size = c - '0'
            if (isFile) {
                usedBlocks += size
                for (p in 1..size) {
                    bitmap[pos++] = fileId
                }
                fileId++
            } else {
                pos += size
            }
            isFile = !isFile
        }
        var backPos = diskSize - 1
        for (i in 0 until usedBlocks) {
            if (bitmap[i] == -1) {
                while (bitmap[backPos] < 0) {
                    backPos--
                }
                bitmap[i] = bitmap[backPos]
                bitmap[backPos--] = -1
            }
        }
        return bitmap.take(usedBlocks).map(Int::toLong).reduceIndexed { index, acc, i -> acc + index * i }
    }

    fun part2(input: List<String>): Long {
        val diskSize = input[0].sumOf { it - '0' }
        val bitmap = IntArray(diskSize) { 0 }
        var pos = 0
        var isFile = true
        var fileId = 0
        val freeList = HashMap<Int, MutableList<Int>>()
        val fileList = ArrayList<Pair<Int, Pair<Int, Int>>>()
        val newFileList = ArrayList<Pair<Int, Pair<Int, Int>>>()
        for (c in input[0]) {
            val size = c - '0'
            if (isFile) {
                if (size == 0) println("Narf!")
                fileList.add(fileId to (pos to size))
                fileId++
            } else {
                if (size > 0) {
                    freeList.getOrPut(size) { ArrayDeque(0) }.add(pos)
                }
            }
            pos += size
            isFile = !isFile
        }

        for (file in fileList.reversed()) {
            val (filePos, size) = file.second
            var bestList: MutableList<Int>? = null
            var bestListSize = 0
            for (s in size..9) {
                val list = freeList[s]
                if (list?.isNotEmpty() == true && list[0] < filePos) {
                    if (bestList == null || list[0] < bestList[0]) {
                        bestList = list
                        bestListSize = s
                    }
                }
            }
            if (bestList != null) {
                val newPos = bestList.removeAt(0)
                if (bestListSize > size) {
                    val remBucket = freeList.getOrPut(bestListSize - size) { ArrayDeque(0) }
                    remBucket.add(newPos + size)
                    remBucket.sort()
                }
                newFileList.add(file.first to (newPos to size))
            } else {
                newFileList.add(file)
            }
        }
        for (file in newFileList) {
            for (i in 0 until file.second.second) {
                bitmap[i + file.second.first] = file.first
            }
        }
        return bitmap.map(Int::toLong).reduceIndexed { index, acc, i -> acc + index * i }
    }

    // test if implementation meets criteria from the description, like:
    val testInput = inlineTestInput.trim().reader().readLines()
    //val testInput = readInput("aoc2024/Day09_test")
    val testInputPart1Result = part1(testInput)
    println("Part 1 Test: $testInputPart1Result")
    val testInputPart2Result = part2(testInput)
    println("Part 2 Test: $testInputPart2Result")
    check(testInputPart1Result == 1928L)
    check(testInputPart2Result == 2858L)

    val input = readInput("aoc2024/Day09")
    part1(input).println()
    part2(input).println()
}