module Part1 where

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


applyModule :: Modules -> Int -> Int -> String -> Bool -> String -> (String, [String], Bool, Modules, Int, Int)
applyModule modules low high source p n
  | notMember n modules = (n, [], True, modules, low, high)
  | otherwise           =
      let m = modules ! n
      in case m of
              Broadcaster o   -> (n, o, p, modules, if not p then low + length o else low,
                                                    if p then high + length o else high)
              FlipFlop s o    -> if not p then let newFF = FlipFlop {state = not s, outputs = o}
                                               in (n, o, not s, insert n newFF modules,
                                                   if s then low + length o else low,
                                                   if not s then high + length o else high)
                                 else (n, [], s, modules, low, high)
              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, low + length o, high)
                                    else (n, o, True, insert n newC modules, low, high + length o)

applyModules :: Modules -> Int -> Int -> String -> Bool -> [String] -> ([(String, [String], Bool)], Modules, Int, Int)
applyModules modules low high _ _ [] = ([], modules, low, high)
applyModules modules low high source p (h: t) =
  let (newS, outputs, newP, newModules, newLow, newHigh) = applyModule modules low high source p h
      (result, newNewModules, newNewLow, newNewHigh) = applyModules newModules newLow newHigh source p t
  in ((newS, outputs, newP): result, newNewModules, newNewLow, newNewHigh)

applySteps :: Modules -> Int -> Int -> [(String, [String], Bool)] -> (Modules, Int, Int)
applySteps modules low high [] = (modules, low, high)
applySteps modules low high ((source, names, pulse): t) =
  let (result, newModules, newLow, newHigh) = applyModules modules low high source pulse names
  in applySteps newModules newLow newHigh (t ++ result)

pressButton :: Modules -> (Modules, Int, Int)
pressButton modules = applySteps modules 1 0 [("button", ["broadcaster"], False)]

pressButtonNTimes' :: Modules -> Int -> Int -> Int -> Int -> (Int, Int)
pressButtonNTimes' modules low high i n
  | i == n    = (low, high)
  | otherwise = let (newModules, newLow, newHigh) = pressButton modules
                in pressButtonNTimes' newModules (low + newLow) (high + newHigh) (i + 1) n

pressButtonNTimes :: Modules -> Int -> (Int, Int)
pressButtonNTimes modules = pressButtonNTimes' modules 0 0 0