package aoc.day5

import aoc._
import scala.collection.immutable.SortedMap
import scala.collection.mutable.ArrayBuffer
import scala.collection.immutable.NumericRange.Exclusive

case class Conversion(source: Long, dest: Long, range: Long):
  val gap = dest - source
  def apply(input: Long) = if source <= input && input < source + range then Some(dest + (input - source))
  else None

  def apply(r: LRange): (Option[LRange], Option[LRange]) =
    assert(r.start >= source)
    val overlap =
      if r.start < source + range then
        Some(
          r.start + gap until (r.`end` min source + range) + gap
        )
      else None
    val rest =
      if r.`end` >= source + range then Some((source + range max r.start) until r.`end`)
      else None
    (overlap, rest)

type LRange = Exclusive[Long]

case class ConvMap(
    val src: String,
    val dest: String,
    conversions: Seq[Conversion]
):
  private val entries = SortedMap(conversions.map(c => (c.source -> c))*)

  def apply(input: Long) =
    entries
      .maxBefore(input + 1) // largest entry with index <= input
      .flatMap((_, c) => c(input))
      .getOrElse(input)

  def apply(input: LRange): Seq[LRange] =
    val buf = ArrayBuffer[LRange]()
    def handleGap(range: LRange)(next: Conversion): Option[LRange] =
      if range.start >= next.source then Some(range) // no gap
      else if range.`end` < next.source then // all in gap
        buf += range
        None
      else
        buf += range.start until next.source
        Some(next.source until range.`end`)
    val endRange = entries.foldLeft[Option[LRange]](Some(input)) {
      case (None, _) => None
      case (Some(range), (_, conv)) =>
        handleGap(range)(conv) match
          case None => None
          case Some(value) =>
            val (overlap, rest) = conv(value)
            // println(s"$src $range => $value $overlap $rest")
            buf ++= overlap
            rest
    }
    (buf ++ endRange).toSeq

object Parser extends CommonParser:
  val nums = rep1(long)

  val seeds = "seeds:" ~> nums

  val mapName = """\w+""".r ~ "-to-" ~ """\w+""".r ~ "map" ^^ { case (src ~ _ ~ dest ~ _) =>
    (src, dest)
  }
  val mapEntry = long ~ long ~ long ^^ { case (dest ~ src ~ range) =>
    Conversion(src, dest, range)
  }
end Parser

val seeds = Parser.parse(Parser.seeds, lines.next()).get

extension [T](xs: Seq[T])
  def splitBy(cond: T => Boolean): Seq[Seq[T]] =
    val splitPos = xs.indexWhere(cond)
    if splitPos == -1 then Seq(xs)
    else xs.take(splitPos) +: xs.drop(splitPos + 1).splitBy(cond)

val maps = lines.toSeq.tail /* remove first empty line */
  .splitBy(_ == "")
  .map { case title :: entries =>
    val (src, dest) = Parser.parse(Parser.mapName, title).get
    src -> ConvMap(src, dest, entries.map(Parser.parse(Parser.mapEntry, _).get))
  }
  .toMap

def toLocation(seed: Long) =
  @scala.annotation.tailrec
  def loop(entryName: String, value: Long): Long =
    if entryName == "location" then value
    else
      val m = maps(entryName)
      loop(m.dest, m(value))
  loop("seed", seed)

def toLocation(seed: LRange) =
  @scala.annotation.tailrec
  def loop(entryName: String, value: Seq[LRange]): Seq[LRange] =
    if entryName == "location" then value
    else
      val m = maps(entryName)
      loop(m.dest, value.flatMap(m(_)))
  loop("seed", Seq(seed))

// Part 1

def part1 =
  val res = seeds.map(toLocation).min
  println(res)

// Part 2

val seedSeqs = seeds.grouped(2).map { case Seq(a, b) => (a until a + b) }

def part2 =
  val res = seedSeqs.map(toLocation).flatMap(_.map(_.start)).min
  println(res)

@main def Day5(part: Long) = part match
  case 1 => part1
  case 2 => part2