import           Data.Bits ((.&.))
import           Data.Char (ord)
import           Data.List (deleteBy, findIndex, foldl')
import           Data.List.Split (splitOn)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import           Data.Vector (Vector, MVector)

type Command = (String, Maybe Int)

getHash :: String -> Int
getHash = foldl' (\acc x -> (acc + ord x) * 17 .&. 255) 0

parseCommand :: String -> Command
parseCommand s
  | last s == '-' = (init s, Nothing)
  | otherwise = (a, Just . read $ b)
  where
    (a:b:_) = splitOn "=" s

sumFocusPower :: Vector (Vector (String, Int)) -> Int
sumFocusPower = timesIndexAndSum . V.map (timesIndexAndSum . V.map snd)
  where
    timesIndexAndSum = V.foldl' (+) 0 . V.imap ((*) . succ)

performCommand
  :: MVector MV.RealWorld (Vector (String, Int)) -> Command -> IO ()
performCommand boxes (label, mFoc) = maybe performDash performEquals mFoc
  where
    boxN = getHash label

    readBox = MV.read boxes boxN

    writeBox = MV.write boxes boxN

    findPred = (== label) . fst

    performDash = writeBox . V.filter (not . findPred) =<< readBox

    performEquals foc = do
      box <- readBox
      let newLens = (label, foc)
          addNewLens = V.snoc box newLens
          swapLens i = V.modify (\v -> MV.write v i newLens) box
      writeBox . maybe addNewLens swapLens . V.findIndex findPred $ box

main :: IO ()
main = do
  commandStrs <- splitOn "," . concat . lines <$> readFile "data.txt"
  boxes <- MV.replicate 256 V.empty
  mapM_ (performCommand boxes . parseCommand) commandStrs
  print . sumFocusPower =<< V.unsafeFreeze boxes