package aoc.day3

import aoc._
import scala.collection.mutable.ArrayBuffer
import scala.math.BigInt

val board = lines.toArray

val numbers = board.toIterator.zipWithIndex.flatMap { (line, row) =>
  var col = 0
  var buf: ArrayBuffer[(Int, Int, String)] = ArrayBuffer.empty
  while col < line.length do
    if line(col).isDigit then
      val numStr = line.slice(col, line.length()).takeWhile(_.isDigit)
      buf += ((row, col, numStr))
      col += numStr.length()
    else col += 1
  buf
}.toSeq

def adjacent(x: Int, y: Int) =
  Seq(-1, 0, 1).flatMap { dx => Seq(-1, 0, 1).map(dy => (x + dx, y + dy)) }

// Part 1

def adjacentCells(x: Int, y: Int) =
  adjacent(x, y)
    .filter((x, y) =>
      x >= 0 && x < board.length && y >= 0 && y < board(x).length
    )
    .map((x, y) => board(x)(y))

def part1 =
  val numbersWithSyms = numbers
    .filter((x, y, num) =>
      (y until y + num.length()).exists(y =>
        adjacentCells(x, y).exists(c => !c.isDigit && c != '.')
      )
    )
    .map(v => BigInt(v._3))
    .sum
  println(numbersWithSyms)

// Part 2

def adjacentNums(x: Int, y: Int) =
  numbers.filter((nx, ny, n) =>
    (x - nx).abs <= 1 && ny <= y + 1 && ny + n.length() >= y
  )

def part2 =
  val gears = board.toIterator.zipWithIndex
    .flatMap { (line, row) =>
      line.zipWithIndex.filter(_._1 == '*').map(v => (row, v._2))
    }
    .map(adjacentNums)
    .filter(_.length == 2)
    .map(s => s.map(v => BigInt(v._3)).product)
    .sum
  println(gears)

@main def Day3(part: Int) = part match
  case 1 => part1
  case 2 => part2