import Data.Char
import Data.List
import Data.List.Split
import Linear.V2

main :: IO ()
main = do
  input <- init . splitOn "\n" <$> readFile "input"
  print $ day3a input
  print $ day3b input

parse :: [String] -> [(V2 Int, Int)]
parse = concat . zipWith (parseLine 0) [0 ..]

parseLine :: Int -> Int -> String -> [(V2 Int, Int)]
parseLine _ _ [] = []
parseLine x y (s:ss)
  | isDigit s =
    (V2 x y, read num) : parseLine (x + length num) y (dropWhile isDigit ss)
  | otherwise = parseLine (succ x) y ss
  where
    num = takeWhile isDigit (s : ss)

getSurround :: (V2 Int, Int) -> [V2 Int]
getSurround (V2 x y, n) =
  line (pred y) ++ line (succ y) ++ [V2 (pred x) y, V2 (x + l) y]
  where
    l = length $ show n
    line c = map (`V2` c) [pred x .. l + x]

checkPart :: (Char -> Bool) -> [String] -> V2 Int -> Bool
checkPart p input (V2 x y)
  | x < 0 || y < 0 = False
  | x >= length (head input) || y >= length input = False
  | otherwise = p (input !! y !! x)

day3a :: [String] -> Int
day3a input =
  sum $
  map snd $
  filter
    (any (checkPart (`notElem` ('.' : ['0' .. '9'])) input) . getSurround)
    (parse input)

day3b :: [String] -> Int
day3b input =
  sum $
  map (product . map (snd . fst)) $
  filter (\x -> length x >= 2) $
  groupBy (\(_, x) (_, y) -> x == y) $ sortOn snd candidates
  where
    candidates =
      [ (x, y)
      | x <- parse input
      , y <- getSurround x
      , _ <- filter (checkPart (== '*') input) [y]
      ]