app "AoC"
    packages {
        pf: "https://github.com/roc-lang/basic-cli/releases/download/0.7.0/bkGby8jb0tmZYsy2hg1E_B2QrCgcSTxdUlHtETwm5m4.tar.br",
        parser: "https://github.com/lukewilliamboswell/roc-parser/releases/download/0.2.0/dJQSsSmorujhiPNIvJKlQoI92RFIG_JQwUfIxZsCSwE.tar.br",
    }
    imports [
        pf.Stdout,
        pf.Task.{ Task },
        parser.Core.{
            Parser,
            const,
            keep,
            skip,
            many,
            oneOf,
        },
        parser.String.{
            digit,
            string,
            parseStr,
            anyCodeunit,
        },
        "day-3-input.txt" as input : Str,
    ]
    provides [main] to pf

main : Task {} *
main =
    Stdout.line part1
    |> Task.await \_ -> Stdout.line part2

part1 : Str
part1 =
    (schematic, potentialParts) =
        input
        |> Str.split "\n"
        |> List.dropIf Str.isEmpty
        |> List.mapWithIndex (\s, i -> parseRow i s)
        |> splitIntoParts

    findPartItems schematic potentialParts
    |> List.sum
    |> Num.toStr

part2 : Str
part2 =
    (schematic, potentialParts) =
        input
        |> Str.split "\n"
        |> List.dropIf Str.isEmpty
        |> List.mapWithIndex (\s, i -> parseRow i s)
        |> splitIntoParts

    findGears schematic potentialParts
    |> List.sum
    |> Num.toStr

Item : [Space, Symbol U8, Number Nat]
Row : List Item
Schematic : List Row

PotentialPart : { row : Nat, col : Nat, len : Nat, part : Nat }

index = \schematic, r, c ->
    List.get schematic r
    |> Result.try \row -> List.get row c

getAdj = \r, c, f ->
    rI = Num.toI64 r
    cI = Num.toI64 c
    List.map [rI - 1, rI, rI + 1] \nr ->
        List.keepOks [cI - 1, cI, cI + 1] \nc ->
            if r >= 0 && c >= 0 then f (Num.toNat nr) (Num.toNat nc) else Err OOB
    |> List.join

findPartItems : Schematic, List PotentialPart -> List Nat
findPartItems = \schematic, potentialParts ->
    getPart = \r, c ->
        when index schematic r c is
            Ok (Symbol _) -> Ok 1
            _ -> Err NotAPart

    goPotentialPart = \{ row, col, len, part } ->
        cols = List.range { start: At col, end: Before (col + len) }
        if
            List.any cols \c -> List.sum (getAdj row c getPart) >= 1
        then
            Ok part
        else
            Err NotAPart

    List.keepOks potentialParts goPotentialPart

getNum = \schematic, r, c ->
    when index schematic r c is
        Ok (Number _) -> Ok (r, c)
        _ -> Err NotANum

findGears : Schematic, List PotentialPart -> List Nat
findGears = \schematic, potentialParts ->
    findPart : Nat, Nat -> PotentialPart
    findPart = \r, c ->
        found =
            List.findFirst potentialParts \{ row, col, len } ->
                row == r && c >= col && c < col + len
        when found is
            Ok part -> part
            Err NotFound -> crash "no part found"

    goGear = \r, c ->
        adjNums = getAdj r c (\r1, c1 -> getNum schematic r1 c1)
        adjParts = List.map adjNums (\(r1, c1) -> findPart r1 c1) |> Set.fromList |> Set.toList

        when adjParts is
            [p1, p2] -> Ok (p1.part * p2.part)
            _ -> Err NotAGear

    goItem = \r, c, i ->
        when i is
            Symbol '*' -> goGear r c
            _ -> Err NotAGear

    goRow = \r, row ->
        List.mapWithIndex row \i, c -> goItem r c i

    List.mapWithIndex schematic \row, r -> goRow r row
    |> List.join
    |> List.keepOks \n -> n

splitIntoParts : List (Row, List PotentialPart) -> (Schematic, List PotentialPart)
splitIntoParts = \l ->
    List.walk l ([], []) \(rows, allPotentialParts), (row, potentialParts) ->
        (List.append rows row, List.concat allPotentialParts potentialParts)

parseRow : Nat, Str -> (Row, List PotentialPart)
parseRow = \row, s ->
    pspace = const Space |> skip (string ".")
    pdigit = const Number |> keep digit
    psymbol = const Symbol |> keep anyCodeunit

    pitem : Parser _ Item
    pitem = oneOf [pspace, pdigit, psymbol]

    findPotentialParts : List PotentialPart, Nat, _, Row -> List PotentialPart
    findPotentialParts = \accum, c, partial, items ->
        # Ideally we'd match on (items, partial) here, but that hits an RC inference bug.
        when items is
            [] ->
                when partial is
                    None -> accum
                    Some { col, len, part } ->
                        List.append accum { row, col, len, part }

            [Number n, .. as rest] ->
                when partial is
                    None ->
                        findPotentialParts accum (c + 1) (Some { col: c, len: 1, part: n }) rest

                    Some { col, len, part } ->
                        findPotentialParts
                            accum
                            (c + 1)
                            (Some { col, len: len + 1, part: part * 10 + n })
                            rest

            [_, .. as rest] ->
                when partial is
                    None ->
                        findPotentialParts accum (c + 1) None rest

                    Some { col, len, part } ->
                        findPotentialParts (List.append accum { row, col, len, part }) (c + 1) None rest

    parsedRow =
        when parseStr (many pitem) s is
            Ok items -> items
            Err (ParsingFailure e) | Err (ParsingIncomplete e) -> crash "bad input: '\(e)'"

    expect List.len parsedRow == Str.countUtf8Bytes s

    potentialParts = findPotentialParts [] 0 None parsedRow

    (parsedRow, potentialParts)

The argument is an anonymous function of type:

    Int Natural -> Bool

But any needs its 2nd argument to be:

    Int Unsigned32 -> Bool

lines = List.range { start: At (number.height - 1), end: At (number.height + 1) }

# With part1, part2 at the top level:
› time ./solutions/day3
real        0m1.744s
user        0m1.702s
sys         0m0.043s

# With both definitions moved into main:
real        0m0.232s
user        0m0.229s
sys         0m0.005s

app "hello"
  packages
    { pf: "https://github.com/roc-lang/basic-cli/releases/download/0.7.0/bkGby8jb0tmZYsy2hg1E_B2QrCgcSTxdUlHtETwm5m4.tar.br"
    , parser: "https://github.com/lukewilliamboswell/roc-parser/releases/download/0.3.0/-e3ebWWmlFPfe9fYrr2z1urfslzygbtQQsl69iH1qzQ.tar.br",
    }
  imports
    [ pf.Stdout
    , pf.Task
    , "./day3-input.txt" as input : Str
    , parser.Core.{ Parser, const, keep, map, oneOrMore }
    , parser.String.{ anyCodeunit, digits, oneOf, parseStr, string }
    ]
  provides [main] to pf


main =
  {} <- Stdout.line "The total sum is \(part1 input)" |> Task.await
  {} <- Stdout.line "The gear ratios are \(part2 input)" |> Task.await
  Task.ok {}

part1 = \schematic ->
  width =
    Str.split schematic "\n"
    |> List.first
    |> Result.withDefault ""
    |> Str.countGraphemes
    |> Num.add 1 # The +1 is for the Newline character
    |> Num.toI16

  identifyParts schematic
  |> appendAdjacentSymbol width
  |> List.keepIf \cell -> when cell.type is
                            Correct -> Bool.true
                            _ -> Bool.false
  |> Set.fromList
  |> Set.toList
  |> List.map \part -> when part.cell is
                          Part num -> Result.withDefault (Str.toU32 num) 0
                          _ -> 0
  |> List.sum
  |> Num.toStr

expect part1 exampleInput == "4361"

part2 = \schematic ->
  width =
    Str.split schematic "\n"
    |> List.first
    |> Result.withDefault ""
    |> Str.countGraphemes
    |> Num.add 1 # The +1 is for the Newline character
    |> Num.toI16

  identifyParts schematic
  |> appendAdjacentSymbol width
  |> List.keepIf \cell -> when cell.type is
                          Gear _ -> Bool.true
                          _ -> Bool.false
  |> Set.fromList
  |> Set.toList
  |> List.map \symbol -> when symbol.type is
                        Gear ratio -> ratio
                        _ -> 0
  |> List.sum
  |> Num.toStr


element : Parser (List U8) CellContent
element =
  oneOf [
    digits |> map \d -> Part (Num.toStr d),
    string "." |> map \_ -> Separator,
    string "\n" |> map \_ -> NewLine,
    string "*" |> map \_ -> Symbol PossibleGear,
    anyCodeunit |> map \_ -> Symbol NotAGear,
  ]

expect parseStr element "467" == Ok (Part "467")
expect parseStr element "." == Ok Separator

schematicParser : Parser (List U8) (List CellContent)
schematicParser =
  const (\i -> i)
  |> keep (oneOrMore element)

expect
  parseStr schematicParser smallInput
  == Ok [
    Part "11",
    Separator,
    Part "2",
    NewLine,
    Symbol NotAGear,
    Separator,
    Separator,
    Separator
  ]

identifyParts = \schematic ->
  parseStr schematicParser schematic
  |> Result.withDefault []
  |> List.walkWithIndex [] \state, cell, index ->
        when cell is
          Part num ->
            List.concat
              state
              (List.repeat
                { id: index
                , cell: cell
                , type: Unknown
                }
                (Str.countGraphemes num))

          _ -> List.append
                  state
                  { id: index
                  , cell: cell
                  , type: Unknown
                  }

expect
  x = identifyParts smallInput
  x == [
    { cell: Part "11", id: 0, type: Unknown },
    { cell: Part "11", id: 0, type: Unknown },
    { cell: Separator, id: 1, type: Unknown },
    { cell: Part "2", id: 2, type: Unknown },
    { cell: NewLine, id: 3, type: Unknown },
    { cell: Symbol NotAGear, id: 4, type: Unknown },
    { cell: Separator, id: 5, type: Unknown },
    { cell: Separator, id: 6, type: Unknown },
    { cell: Separator, id: 7, type: Unknown }]

CellType : [Unknown, Correct, Extra, Gear U32, NotAPart]
CellContent : [NewLine,
  Part Str,
  Separator,
  Symbol [PossibleGear, NotAGear],
  ]

appendAdjacentSymbol : List {
  cell : CellContent,
  type: CellType,
  id : Nat,
}, I16 -> List {
  cell : CellContent,
  type : CellType,
  id : Nat,
}
appendAdjacentSymbol = \identifiedParts, width ->
  # width should include the newline
  adjacentCellIndexes = [
    -width - 1,
    -width,
    -width + 1,
    -1,
    1,
    width - 1,
    width,
    width + 1
  ]

  State : List {
      cell : CellContent,
      type : CellType,
      id : Nat,
    }
  Cell : {
      cell : CellContent,
      type: CellType,
      id : Nat,
    }

  x : State, Cell, Nat -> State
  x = \state, cell, index ->
    when cell.cell is
      Part _ ->
        adjacentSymbols =
          List.map adjacentCellIndexes \relativeIndex -> (Num.add relativeIndex (Num.toI16 index))
          |> List.keepIf \absoluteIndex -> absoluteIndex >= 0
          |> List.map Num.toNat
          |> List.keepOks \absoluteIndex -> List.get identifiedParts absoluteIndex
          |> List.keepIf \adjacentCell ->
                when adjacentCell.cell is
                  Symbol _ -> Bool.true
                  _ -> Bool.false

        List.append state {
          cell &
          type: if List.isEmpty adjacentSymbols then Extra else Correct
        }

      Symbol symbolType ->
        when symbolType is
          PossibleGear ->
            adjacentParts =
              List.map adjacentCellIndexes \relativeIndex -> (Num.add relativeIndex (Num.toI16 index))
              |> List.keepIf \absoluteIndex -> absoluteIndex >= 0
              |> List.map Num.toNat
              |> List.keepOks \absoluteIndex -> List.get identifiedParts absoluteIndex
              |> List.keepIf \adjacentCell ->
                    when adjacentCell.cell is
                      Part _ -> Bool.true
                      _ -> Bool.false
              |> Set.fromList
              |> Set.toList

            if List.len adjacentParts >= 2 then
              # _ = debug index "index"
              # _ = debug cell "cell"
              # _ = debug adjacentParts "adjacentParts"

              gearRatio =
                adjacentParts
                |> Set.fromList
                |> Set.toList
                |> List.map \adjacentCell -> when adjacentCell.cell is
                                                Part num -> Result.withDefault (Str.toU32 num) 1
                                                _ -> 1
                |> List.product

              List.append state { cell & type: Gear gearRatio }
            else
              List.append state { cell & type: NotAPart }

          _ -> List.append state { cell & type: NotAPart }

      _ -> List.append state { cell & type: NotAPart }

  List.walkWithIndex identifiedParts [] x

expect identifyParts smallInput |> appendAdjacentSymbol 5 == [
    { cell: Part "11", id: 0, type: Correct },
    { cell: Part "11", id: 0, type: Correct },
    { cell: Separator, id: 1, type: NotAPart },
    { cell: Part "2", id: 2, type: Extra },
    { cell: NewLine, id: 3, type: NotAPart },
    { cell: Symbol NotAGear, id: 4, type: NotAPart },
    { cell: Separator, id: 5, type: NotAPart },
    { cell: Separator, id: 6, type: NotAPart },
    { cell: Separator, id: 7, type: NotAPart }
]



smallInput =
  """
  11.2
  $...
  """ |> Str.trim

smallGearInput =
  """
  11.2
  $.*.
  """ |> Str.trim

exampleInput =
  """
  467..114..
  ...*......
  ..35..633.
  ......#...
  617*......
  .....+.58.
  ..592.....
  ......755.
  ...$.*....
  .664.598..
  """
  |> Str.trim


debug = \v, ctx ->
  dbg ctx
  dbg v
  v

app "AoC"
    packages { pf: "https://github.com/roc-lang/basic-cli/releases/download/0.7.0/bkGby8jb0tmZYsy2hg1E_B2QrCgcSTxdUlHtETwm5m4.tar.br" }

    imports [pf.Stdout, pf.Task.{ Task }, "day-3-input.txt" as input : Str]
    provides [main] to pf

main : Task {} *
main =
    dbg
        P1 (part1 input)

    dbg
        P2 (part2 input)

    Stdout.line "done"

part1 = \str ->
    s = Str.toUtf8 str
    lineLength =
        List.findFirstIndex s (\c -> c == '\n')
        |> unwrap
        |> Num.add 1

    startState = {
        row: 0,
        sum: 0,
        currentNumber: 0,
        shouldAdd: Bool.false,

    }
    s
    |> List.walkWithIndex startState \state, elem, index -> sumPartNumbers state elem s index lineLength

part2 = \str ->
    s = Str.toUtf8 str
    lineLength =
        List.findFirstIndex s (\c -> c == '\n')
        |> unwrap
        |> Num.add 1

    startState = 0
    List.walkWithIndex s startState \state, elem, index -> sumGearRatio state elem s index lineLength

sumGearRatio = \state, elem, inp, index, lineLength ->
    if elem == '*' then
        nearNumbers = getNumbersNear inp index lineLength
        if List.len nearNumbers == 2 then
            state + List.product nearNumbers
        else
            state
    else
        state

getNumbersNear = \list, index, lineLength ->
    getToRight = \pos ->
        if isDigitAt list pos then
            List.dropFirst list pos
            |> takeWhile isDigit
            |> Str.fromUtf8
            |> unwrap
            |> Str.toNat
            |> unwrap
            |> Ok
        else
            Err NoNumber
    getToLeft = \pos ->
        if isDigitAt list pos then
            List.takeFirst list (pos + 1)
            |> takeWhileBackwards isDigit
            |> Str.fromUtf8
            |> unwrap
            |> Str.toNat
            |> unwrap
            |> Ok
        else
            Err NoNumber

    right = getToRight (index + 1)
    left = getToLeft (index - 1)

    topindex = Num.subWrap index lineLength
    # extracting this as a closure creates a compiler error, that's why it is copied for bottom
    top =
        if isDigitAt list topindex then
            # only one possible
            leftPart =
                List.takeFirst list topindex
                |> takeWhileBackwards isDigit
            rightPart =
                List.dropFirst list topindex
                |> takeWhile isDigit
            List.concat leftPart rightPart
            |> Str.fromUtf8
            |> unwrap
            |> Str.toNat
            |> unwrap
            |> List.single
        else
            [getToRight (topindex + 1), getToLeft (topindex - 1)]
            |> List.keepOks \x -> x

    bottomindex = index + lineLength
    bottom =
        if isDigitAt list bottomindex then
            # only one possible
            leftPart =
                List.takeFirst list bottomindex
                |> takeWhileBackwards isDigit
            rightPart =
                List.dropFirst list bottomindex
                |> takeWhile isDigit
            List.concat leftPart rightPart
            |> Str.fromUtf8
            |> unwrap
            |> Str.toNat
            |> unwrap
            |> List.single
        else
            [getToRight (bottomindex + 1), getToLeft (bottomindex - 1)]
            |> List.keepOks \x -> x
    [right, left]
    |> List.keepOks \x -> x
    |> List.concat top
    |> List.concat bottom

takeWhile = \list, cond ->
    when List.findFirstIndex list \elem -> !(cond elem) is
        Ok firstWrongIndex -> List.takeFirst list firstWrongIndex
        Err NotFound -> list

takeWhileBackwards = \list, cond ->
    when List.findLastIndex list \elem -> !(cond elem) is
        Ok lastWrongIndex -> List.dropFirst list (lastWrongIndex + 1)
        Err NotFound -> list

isDigitAt = \list, index ->
    when List.get list index is
        Err _ -> Bool.false
        Ok d -> isDigit d

sumPartNumbers = \state, elem, inp, index, lineLength ->
    if elem == '\n' then
        if state.shouldAdd then
            newState = addToSum state
            { newState & row: state.row + 1 }
        else
            { state & row: state.row + 1, currentNumber: 0 }
    else if isDigit elem then
        { state & currentNumber: state.currentNumber * 10 + (Num.toNat elem) - '0', shouldAdd: state.shouldAdd || hasSymbolNear inp index lineLength }
    else if state.shouldAdd then
        addToSum state
    else
        { state & currentNumber: 0 }

hasSymbolNear = \list, index, lineLength ->
    isSymbolAt list (index + 1)
    || isSymbolAt list (Num.subWrap index 1)
    || isSymbolAt list (index + lineLength)
    || isSymbolAt list (index + lineLength - 1)
    || isSymbolAt list (index + lineLength + 1)
    || isSymbolAt list (Num.subWrap index lineLength)
    || isSymbolAt list (Num.subWrap index (lineLength - 1))
    || isSymbolAt list (Num.subWrap index (lineLength + 1))

isSymbolAt = \list, index ->
    when List.get list index is
        Err _ -> Bool.false
        Ok d -> isSymbol d

isSymbol = \c ->
    (c < '0' || c > '9') && c != '.' && c != '\n'

isDigit = \c ->
    c >= '0' && c <= '9'

addToSum = \state ->
    { state & sum: state.sum + state.currentNumber, shouldAdd: Bool.false, currentNumber: 0 }

example : Str
example =
    """
    467..114..
    ...*......
    ..35..633.
    ......#...
    617*......
    .....+.58.
    ..592.....
    ......755.
    ...$.*....
    .664.598..
    """

unwrap = \res ->
    when res is
        Ok x -> x
        Err _ -> crash "bad unwrap"