module Main where

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

-- TODO: solve part 2 by transposing the matrix before parsing the
-- numbers.

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
  where
    -- 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

    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

data Oper = Mult | Add
  deriving (Show)

part2 :: String -> Int
part2 = (\g -> foldl' (+) 0 $ map (calculateGroup g) [0 .. M.size (columns g) - 1]) . newGrid

-- Extract the character in the grid on row, group and column.
--
-- Return Nothing if the char is ' '.
idx :: Grid -> Int -> Int -> Int -> Maybe Char
idx (Grid cols rows cs) row col n
  | col >= M.size cols = error $ "idx: requested a too large number group index (" ++ show col ++ " >= " ++ show cols ++ ")"
  | row >= rows = error $ "idx: requested a too large row index (" ++ show row ++ " >= " ++ show rows ++ ")"
  | n >= size (cols M.! col) = error $ "idx: requested a too large column index (" ++ show n ++ " >= " ++ show (size (cols M.! col)) ++ ")"
  | otherwise = case cs !! row !! i of
      ' ' -> Nothing
      c -> Just c
  where
    i = n + (sum . map ((+ 1) . (size . (cols M.!))) $ [0 .. col - 1])

digit :: Grid -> Int -> Int -> Int -> Maybe Int
digit g r c = fmap digitToInt . idx g r c

operator :: Grid -> Int -> Oper
operator g col = case idx g (rows g - 1) col 0 of
  Just '+' -> Add
  Just '*' -> Mult
  Just ch -> error $ "operator: Expected '+' or '*', got " <> show ch
  Nothing -> error $ "operator: Expected '+' or '*', got nothing"

-- Calculate a column by traversing all rows at that column inside that group.
--
-- Calculation is performed by first multiplying the accumulated value
-- by 10 and then adding the new digit. If no value exists in that
-- slot we don't modify the accumulated value.
calculateColumn :: Grid -> Int -> Int -> Int
calculateColumn g c n = foldl' (\acc r -> maybe acc ((+) (acc * 10)) $ digit g r c n) 0 [0 .. rows g - 2]

-- Calculates each column and folds them using the operator at the last line.
calculateGroup :: Grid -> Int -> Int
calculateGroup g c = foldl' (op $ oper column) (identity column) $ map (calculateColumn g c) [0 .. (size column) - 1]
  where
    column = columns g M.! c

op :: Oper -> (Int -> Int -> Int)
op Mult = (*)
op Add = (+)

data Grid = Grid
  { columns :: IntMap Column,
    rows :: Int,
    cells :: [[Char]]
  }
  deriving (Show)

data Column = Column
  { size :: Int,
    oper :: Oper,
    identity :: Int
  }
  deriving (Show)

-- | Construct a grid based on the last line.
--
-- Using the last line we can get the number of digits in each column.
newGrid :: String -> Grid
newGrid = go . lines
  where
    go :: [String] -> Grid
    go ls =
      let columns = M.fromList $ foldl' go' [] (last ls)
       in Grid columns (length ls) ls

    go' :: [(Int, Column)] -> Char -> [(Int, Column)]
    go' [] ch = pure . (0,) . (\(oper, idty) -> Column 1 oper idty) $ parseOper ch
    go' ((n, c) : lengths) ch
      | ch == ' ' = (n, c {size = size c + 1}) : lengths
      | otherwise = (n + 1, (\(oper, idty) -> Column 1 oper idty) (parseOper ch)) : (n, c {size = size c - 1}) : lengths

    parseOper :: Char -> (Oper, Int)
    parseOper '*' = (Mult, 1)
    parseOper '+' = (Add, 0)
    parseOper c = error $ "parseOper: expects a '+' or '*', got " <> show c

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