{-# LANGUAGE OverloadedStrings #-}
module Day15 (parse, solveA, solveB) where

import Data.Void (Void)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Char (ord, isSpace)
import Data.List (findIndex)
import Data.Bifunctor (bimap)
import Text.Megaparsec (Parsec, errorBundlePretty, runParser, sepBy, optional, eof, many, choice)
import Text.Megaparsec.Char (newline, alphaNumChar, char)
import Text.Megaparsec.Char.Lexer (decimal)
import Lens.Micro (ix, (%~), (.~))


data Op
  = Set Text Int
  | Del Text
  deriving Show

type Entry = (Text, Int)

type Parser = Parsec Void Text


parse :: Text -> Either String [(Text, Op)]
parse text = bimap errorBundlePretty (zip . simple $ text) . runParser complex "" $ text
  where
    simple :: Text -> [Text]
    simple = T.splitOn "," . T.filter (not . isSpace)
    single :: Parser Op
    single = do
      label  <- T.pack <$> many alphaNumChar
      choice
        [ Set label <$> (char '=' *> decimal)
        , Del label <$  char '-'
        ]
    complex :: Parser [Op]
    complex = single `sepBy` "," <* optional newline <* eof

hash :: Text -> Int
hash = T.foldl step 0
  where
    step :: Int -> Char -> Int
    step curr = (`rem` 256) . (17 *) . (curr +) . ord

solveA :: [(Text, Op)] -> Int
solveA = sum . map (hash . fst)

del :: Text -> [[Entry]] -> [[Entry]]
del k = ix (hash k) %~ filter ((k /=) . fst)

set :: Text -> Int -> [[Entry]] -> [[Entry]]
set k v = ix (hash k) %~ upsert
  where
    update :: Int -> [Entry] -> [Entry]
    update n = ix n .~ (k, v)
    insert :: [Entry] -> [Entry]
    insert = (<> [(k, v)])
    upsert :: [Entry] -> [Entry]
    upsert es = maybe (insert es) (`update` es) . findIndex ((k ==) . fst) $ es

solveB :: [(Text, Op)] -> Int
solveB = sum . zipWith row [1..] . foldl (flip run) (replicate 256 []) . map snd
  where
    run :: Op -> [[Entry]] -> [[Entry]]
    run op = case op of
      Set k v -> set k v
      Del k   -> del k
    box :: Int -> Entry -> Int
    box i = (i *) . snd
    row :: Int -> [Entry] -> Int
    row i = (i *) . sum . zipWith box [1..]