module Main (main) where

import Data.List (elemIndex)
import Data.Maybe (fromJust)
import Lib.Parser (Parser (parse), char, many1, number, separated)

-- | Part 1
newtype Rule = Rule (Int, Int) deriving (Show)

newtype Update = Update [Int] deriving (Show)

rule :: Parser Rule
rule = (\a _ b -> Rule (a, b)) <$> number <*> char '|' <*> number

update :: Parser Update
update = Update <$> char ',' `separated` number

safetyManual :: Parser ([Rule], [Update])
safetyManual = filterParsed <$> many1 rule' <*> char '\n' <*> many1 update'
  where
    filterParsed :: [Rule] -> Char -> [Update] -> ([Rule], [Update])
    filterParsed r _ u = (r, u)

    rule' :: Parser Rule
    rule' = const <$> rule <*> char '\n'

    update' :: Parser Update
    update' = const <$> update <*> char '\n'

validateRule :: Update -> Rule -> Bool
validateRule (Update u) (Rule (a, b)) = validateRule' $ (<) <$> elemIndex a u <*> elemIndex b u
  where
    validateRule' :: Maybe Bool -> Bool
    validateRule' Nothing = True
    validateRule' (Just result) = result

validateRules :: [Rule] -> Update -> Bool
validateRules rules upd = all (validateRule upd) rules

middle :: Update -> Int
middle (Update u) = u !! (length u `div` 2)

validUpdates :: ([Rule], [Update]) -> [Update]
validUpdates (r, u) = filter (validateRules r) u

part1 :: String -> IO ()
part1 s = print $ sum $ map middle $ validUpdates $ fst $ fromJust $ parse safetyManual s

-- | Part 2
invalidUpdates :: ([Rule], [Update]) -> [Update]
invalidUpdates (r, u) = filter (not . validateRules r) u

insertPage :: [Rule] -> Int -> Update -> Update
insertPage rules page (Update pages) = insertPage' rules pages page 0
  where
    insertPage' :: [Rule] -> [Int] -> Int -> Int -> Update
    insertPage' [] pages' page' minViable = let (ps,ps') = splitAt minViable pages' in Update $ ps ++ [page'] ++ ps'
    insertPage' _ [] page' _ = Update [page']
    insertPage' ((Rule (a, b)):rs) pages' page' minViable
      | page' == b && a `elem` pages' = insertPage' rs pages' page' (max minViable (((+1) . fromJust . elemIndex a) pages'))
      | otherwise = insertPage' rs pages page minViable

fixUpdate :: [Rule] -> Update -> Update
fixUpdate rules (Update pages) = foldr (insertPage rules) (Update []) pages

part2 :: String -> IO ()
part2 s =
  let (rules, updates) = (fst . fromJust . parse safetyManual) s
      invUpd = invalidUpdates (rules, updates)
  in print $ sum $ map (middle . fixUpdate rules) invUpd

-- | Main
main :: IO ()
main = readFile "resources/day5.txt" >>= part2