module Part2 where

import Commons
import Data.Map (insert, foldr, (!), notMember, toList, keys)


applyModule :: Modules -> String -> String -> Bool -> String -> (String, [String], Bool, Modules, Bool)
applyModule modules final source p n
  | notMember n modules = (n, [], True, modules, False)
  | otherwise           =
      let m = modules ! n
          partFinal = n == final && p
      in case m of
              Broadcaster o   -> (n, o, p, modules, partFinal)
              FlipFlop s o    -> if not p then let newFF = FlipFlop {state = not s, outputs = o}
                                               in (n, o, not s, insert n newFF modules, partFinal)
                                 else (n, [], s, modules, partFinal)
              Conjonction i o -> let newC = Conjonction {inputs = insert source p i, outputs = o}
                                     state = Data.Map.foldr (&&) True $ inputs newC
                                 in if state then (n, o, False, insert n newC modules, n == final)
                                    else (n, o, True, insert n newC modules, partFinal)

applyModules :: Modules -> String -> String -> Bool -> [String] -> ([(String, [String], Bool)], Modules, Bool)
applyModules modules _ _ _ [] = ([], modules, False)
applyModules modules final source p (h: t) =
  let (newS, outputs, newP, newModules, done) = applyModule modules final source p h
      (result, newNewModules, newDone) = applyModules newModules final source p t
  in ((newS, outputs, newP): result, newNewModules, done || newDone)

applySteps :: Modules -> String -> [(String, [String], Bool)] -> (Modules, Bool)
applySteps modules _ [] = (modules, False)
applySteps modules final ((source, names, pulse): t) =
  let (result, newModules, done) = applyModules modules final source pulse names
      (newNewModules, newDone) = applySteps newModules final (t ++ result)
  in (newNewModules, done || newDone)

pressButton :: Modules -> String -> (Modules, Bool)
pressButton modules final = applySteps modules final [("button", ["broadcaster"], False)]

getFinalConjonction :: Modules -> String
getFinalConjonction = fst . head . filter (\ (_, v) -> "rx" `elem` outputs v) . toList

pressButtonUntilDone' :: Modules -> String -> Int -> [Int] -> [Int]
pressButtonUntilDone' modules final i nPresses
  | length nPresses == length (inputs (modules ! final)) = nPresses
  | otherwise                                            =
      let (newModules, done) = pressButton modules final
      in pressButtonUntilDone' newModules final (i + 1) (if done then (i + 1): nPresses else nPresses)

pressButtonUntilDone :: Modules -> [Int]
pressButtonUntilDone modules = pressButtonUntilDone' modules (getFinalConjonction modules) 0 []