import scala.io.Source

object Day3 extends App:

  val input = Source
    .fromURL(getClass.getResource("day3Input.txt"))
    .mkString
    .split('\n')
    .map(_.toArray.map(Bit.parse))

  enum Bit:
    case One, Zero

    def flip() = this match
      case One  => Zero
      case Zero => One

    override def toString() = this match
      case One  => "1"
      case Zero => "0"

  object Bit:
    def parse(char: Char): Bit = char match
      case '1' => One
      case '0' => Zero

  import Bit._

  def findMostCommonBit(bits: Array[Bit]): Bit =
    if (bits.count(_ == One) > bits.count(_ == Zero)) One
    else Zero

  val gamma = input.transpose
    .map(findMostCommonBit)

  val epsilon = gamma.map(_.flip())

  val part1 =
    Integer.parseInt(gamma.mkString, 2) * Integer.parseInt(epsilon.mkString, 2)

  println(s"Part 1: $part1")

  // Part 2

  enum BitCriteria(val selectBitByCount: (Int, Int) => Bit):
    case OxygenGeneratorRating
        extends BitCriteria((onesCount, zerosCount) =>
          if (onesCount >= zerosCount) One else Zero
        )
    case CO2ScrubberRating
        extends BitCriteria((onesCount, zerosCount) =>
          if (onesCount < zerosCount) One else Zero
        )

  import BitCriteria._

  def filterByBitCriteria(
      filteredBitIndex: Int,
      numbers: Array[Array[Bit]],
      bitCriteria: BitCriteria
  ): Array[Bit] =
    val filteredBits = numbers.map(_(filteredBitIndex))
    val onesCount = filteredBits.count(_ == One)
    val zeroesCount = filteredBits.count(_ == Zero)
    val headToKeep = bitCriteria.selectBitByCount(onesCount, zeroesCount)
    val filteredNumbers = numbers.filter(_(filteredBitIndex) == headToKeep)

    if (filteredNumbers.length == 1)
      filteredNumbers.head
    else
      filterByBitCriteria(filteredBitIndex + 1, filteredNumbers, bitCriteria)

  val oxygenGeneratorRating =
    Integer.parseInt(
      filterByBitCriteria(0, input, OxygenGeneratorRating).mkString,
      2
    )
  val cO2ScrubberRating =
    Integer.parseInt(
      filterByBitCriteria(0, input, CO2ScrubberRating).mkString,
      2
    )

  println(s"Part 2: ${oxygenGeneratorRating * cO2ScrubberRating}")