module Day3

import Prelude
import Node

pfunc repr : {a : U} -> a -> String := `(a,o) => ''+o`
pfunc jrepr : {a : U} -> a -> String := `(a,o) => JSON.stringify(o, null, '  ')`


maybe : ∀ a b. b → (a → b) → Maybe a → b
maybe def f Nothing = def
maybe def f (Just a) = f a

-- was 'structure' I could make a `record` that destructures to this..
data Number : U where
  MkNumber : (start : Nat) -> (stop : Nat) → (value : Int) → Number

isDigit : Char -> Bool
isDigit '0' = True
isDigit '1' = True
isDigit '2' = True
isDigit '3' = True
isDigit '4' = True
isDigit '5' = True
isDigit '6' = True
isDigit '7' = True
isDigit '8' = True
isDigit '9' = True
isDigit _ = False

numbers : List Char -> List Number
numbers arr = go arr Z
  where
    go : List Char → Nat → List Number
    go (c :: cs) start = if isDigit c
      -- then let (front,back) = span isDigit (c :: cs) in ?
      then case span isDigit (c :: cs) of
        -- NOW FC on app is now broken, need the fc of the left
          (front,back) => let stop = start + length front
            in MkNumber start stop (stringToInt $ pack front) :: go back stop
      else go cs (S start)
    go Nil start = Nil


range : ∀ a. Nat -> Nat -> List a -> List a
range _ _ Nil = Nil
range _ Z _ = Nil
range Z (S k) (x :: xs) = x :: range Z k xs
range (S n) (S m) (x :: xs) = range n m xs

isPart : List (List Char) -> Nat -> Number -> Bool
isPart rows row (MkNumber start end _) =
  checkRow (pred row) || checkRow row || checkRow (S row)
  where
    isThing : Char -> Bool
    isThing c = not (isDigit c || c == '.')

    checkRow : Nat -> Bool
    checkRow r = case getAt r rows of
      Nothing => False
      Just chars => case filter isThing (range (pred start) (S end) chars) of
        Nil => False
        _ => True

getValue : Number -> Int
getValue (MkNumber _ _ v) = v

part1 : List (List Char) -> Int
part1 rows =
    foldl (\ acc num => acc + getValue num) 0 $
      join $ map (partNums rows) $ enumerate rows
  where
    partNums : List (List Char) -> (Nat × List Char) -> List Number
    partNums grid (r, cs) =
      filter (isPart grid r) $ (numbers cs)

gears : List (List Char) -> List Char -> Nat -> Int
gears rows row y =
    let a = numbers (getAt! (pred y) rows)
        b = numbers (getAt! y rows)
        c = numbers (getAt! (S y) rows)
        all = a ++ b ++ c
        cands = map fst $ filter (_==_ '*' ∘ snd) (enumerate row)
    in foldl _+_ 0 $ map (check all) cands
  where
    ratio : List Int → Int
    ratio (a :: b :: Nil) = a * b
    ratio _ = 0

    match : Nat → Number → Bool
    match y (MkNumber start stop value) = pred start <= y && y < S stop

    check : List Number → Nat →  Int
    check nums y = ratio $ map getValue (filter (match y) nums)

part2 : List (List Char) -> Int
part2 rows =
  foldl go 0 (enumerate rows)
  where
    go : Int → Nat × List Char → Int
    go acc (y, row) = acc + gears rows row y

-- 4361 / 467835
-- 517021 / 81296995
run : String -> IO Unit
run fn = do
  content <- readFile fn
  let grid = (splitOn '\n' $ unpack $ trim content)
  putStrLn fn
  printLn (part1 grid)
  printLn (part2 grid)

main : IO Unit
main = do
  run "aoc2023/day3/eg.txt"
  run "aoc2023/day3/input.txt"