module Day18 where
import Data.Char (isDigit, digitToInt)
import Data.List (foldl')

data Number = Single Int | Pair Number Number deriving (Eq)

instance Show Number where
  show (Single x) = show x
  show (Pair x y) = '[':show x ++ ',':show y ++ "]"

replaceRight :: Int -> Number -> Number
replaceRight y (Single x) = Single (x+y)
replaceRight y (Pair l r) = Pair l (replaceRight y r)

replaceLeft :: Int -> Number -> Number
replaceLeft y (Single x) = Single (x+y)
replaceLeft y (Pair l r) = Pair (replaceLeft y l) r

explode :: Int -> Number -> (Maybe Int, Maybe Int, Number)
explode _ (Single x) = (Nothing, Nothing, Single x)
explode 4 (Pair (Single x) (Single y)) = (Just x, Just y, Single 0)
explode 4 _ = error "pairs that are deeper than 4"
explode lvl (Pair x y)
        | ln /= x = case lr of
                Nothing  -> (ll, Nothing, Pair ln y)
                (Just r) -> (ll, Nothing, Pair ln (replaceLeft r y))
        | otherwise = case rl of
                Nothing  -> (Nothing, rr, Pair x rn)
                (Just l) -> (Nothing, rr, Pair (replaceRight l x) rn)
        where
                (ll, lr, ln) = explode (lvl+1) x
                (rl, rr, rn) = explode (lvl+1) y

explodeFull :: Number -> Number
explodeFull n = let (_,_,new) = explode 0 n in new

split :: Number -> Number
split (Single x)
        | x > 9 = Pair (Single $ x `div` 2) (Single $ (x+1) `div` 2)
        | otherwise = Single x
split (Pair x y) = if split x /= x then Pair (split x) y else Pair x (split y)

reduce :: Number -> (Number, Bool)
reduce n
  | n /= exploded = (exploded, True)
  | n /= splitN = (splitN, True)
  | otherwise = (n, False)
  where
      exploded = explodeFull n
      splitN = split n

reduceFull :: Number -> Number
reduceFull n = fst $ until (not . snd) (reduce . fst) (n, True)

parse :: String -> (Number, String)
parse [] = error "parsing empty string"
parse ('[': xs) = (Pair l r, xs'') where
                (l, ',':xs') = parse xs
                (r, ']':xs'') = parse xs'
parse (x:xs)
        | isDigit x = (Single $ digitToInt x, xs) -- yes i know this only works for single-digit numbers
        | otherwise = error "parse error"

add :: Number -> Number -> Number
add = Pair

addAll :: [Number] -> Number
addAll [] = error "sum of empty list"
addAll (x:xs) = foldl' (\s i -> reduceFull (s `add` i)) x xs

magnitude :: Number -> Int
magnitude (Single x) = x
magnitude (Pair l r) = (magnitude l * 3) + (magnitude r * 2)

main :: IO ()
main = do
        nums <- map (fst . parse) . lines <$> getContents
        print $ addAll nums

        putStr "part 1: "
        print $ magnitude $ addAll nums

        putStr "part 2: " -- please use compiled code!
        let m = maximum $ map (\n1 -> maximum $ map (\n2 ->
                magnitude $ reduceFull (n1 `add` n2)) nums) nums
        print m