module Main where

import Data.Bifunctor (first, second)
import Data.Bits
import Numeric.Natural

import Input


data Packet = Lit Int Natural | Op Int Int [Packet]
  deriving (Show)

dec :: [Bool] -> Natural
dec l = sum (zipWith (*) (map (fromIntegral . fromEnum) (reverse l)) (iterate (*2) 1))

parse :: [Bool] -> (Packet, [Bool])
parse l =
    let (ver, l1) = splitAt 3 l
        (typ, l2) = splitAt 3 l1
    in case dec typ of
         4 -> let collectGroups m =
                    let (g1, m1) = splitAt 5 m
                    in case g1 of False : g1' -> (g1', m1)
                                  True : g1' -> first (g1' ++) (collectGroups m1)
                                  [] -> error "Unexpected EOF in literal"
              in first (Lit (fromIntegral (dec ver)) . dec) (collectGroups l2)
         typ' | False : l3 <- l2 ->
                  let (bodylen, l4) = splitAt 15 l3
                      (body, l5) = splitAt (fromIntegral (dec bodylen)) l4
                  in (Op (fromIntegral (dec ver)) (fromIntegral typ') (parseMany body), l5)
              | True : l3 <- l2 ->
                  let (numpkt, l4) = splitAt 11 l3
                      (pkts, l5) = parseN (dec numpkt) l4
                  in (Op (fromIntegral (dec ver)) (fromIntegral typ') pkts, l5)
              | otherwise -> error "Unexpected EOF before length ID bit in operator"

parseMany :: [Bool] -> [Packet]
parseMany [] = []
parseMany l = uncurry (:) (second parseMany (parse l))

parseN :: Natural -> [Bool] -> ([Packet], [Bool])
parseN 0 l = ([], l)
parseN n l = let (pkt, rest) = parse l
                 (pkts, rest') = parseN (n - 1) rest
             in (pkt : pkts, rest')

operator :: Int -> [Natural] -> Natural
operator 0 = sum
operator 1 = product
operator 2 = minimum
operator 3 = maximum
operator 5 = \[a,b] -> fromIntegral (fromEnum (a > b))
operator 6 = \[a,b] -> fromIntegral (fromEnum (a < b))
operator 7 = \[a,b] -> fromIntegral (fromEnum (a == b))
operator n = error ("Invalid operator " ++ show n)

eval :: Packet -> Natural
eval (Lit _ n) = n
eval (Op _ n ps) = operator n (map eval ps)

main :: IO ()
main = do
    let decodeHex c | '0' <= c, c <= '9' = fromEnum c - fromEnum '0'
                    | 'A' <= c, c <= 'F' = 10 + fromEnum c - fromEnum 'A'
                    | otherwise = error "Invalid hex digit"
        toBinary4 n = [testBit n i | i <- [3,2,1,0]]
    bits <- concatMap (toBinary4 . decodeHex) . head <$> getInput 16
    let addVers (Lit v _) = v
        addVers (Op v _ ps) = v + sum (map addVers ps)
    print (addVers (fst (parse bits)))
    print (eval (fst (parse bits)))