module Main where

import Data.Either (lefts, rights)
import Data.IntMap (IntMap)
import qualified Data.IntMap as M
import Data.List (foldl')
import System.Environment (getArgs)

main :: IO ()
main = do
  input <- readFile . head =<< getArgs
  putStrLn $ unwords ["Part 1:", show $ part1 input]
  putStrLn $ unwords ["Part 2:", show $ part2 input]

part1 :: String -> Int
part1 = M.foldl' (+) 0 . parse

data Oper = Mult | Add
  deriving (Show)

-- Parse each space-delimited value as either a number or an
-- operator. Then combine the operator and the operands into a
-- calculation.
parse :: String -> IntMap Int
parse = M.map complete . foldl' insertPartial mempty . concat . map parseLine . lines
  where
    parseLine :: String -> [(Int, Either Oper Int)]
    parseLine = zip [0 ..] . map parseWord . words

    parseWord :: String -> Either Oper Int
    parseWord "*" = Left Mult
    parseWord "+" = Left Add
    parseWord w = pure $ read w

    insertPartial :: IntMap [Either Oper Int] -> (Int, Either Oper Int) -> IntMap [Either Oper Int]
    insertPartial m (column, oper) = M.insertWith ((<>)) column [oper] m

    complete :: [Either Oper Int] -> Int
    complete m = case (rights m, head (lefts m)) of
      (nums, Mult) -> product nums
      (nums, Add) -> sum nums

part2 :: String -> Int
part2 = error "Not implemented"

testInput :: String
testInput =
  unlines
    [ "123 328  51 64",
      " 45 64  387 23",
      "  6 98  215 314",
      "*   +   *   + "
    ]