module Main (main) where

import Control.Monad (replicateM)
import Control.Parallel.Strategies (parList, rpar, using)
import Data.Maybe (fromJust, isJust)
import Lib.Parser (Parser (parse), char, many1, number, separated, string)

-- | Part 1
data Operator = Multiplication | Addition | Concatenation | Number Int deriving (Show, Eq)

newtype Equation = Equation (Int, [Int]) deriving (Show)

equation :: Parser Equation
equation = (\equationResult _ equationParts -> Equation (equationResult, equationParts)) <$> number <*> string ": " <*> (char ' ' `separated` number)

checkEquation :: Equation -> [Operator] -> Bool
checkEquation (Equation (_, [])) _ = error "Empty equation!"
checkEquation (Equation (left, r : rs)) ops = checkEquation' rs ops r
  where
    checkEquation' :: [Int] -> [Operator] -> Int -> Bool
    checkEquation' [] [] n = left == n
    checkEquation' [] _ _ = error "Have more operators than numbers!"
    checkEquation' _ [] _ = error "Have more numbers than operators!"
    checkEquation' (n : ns) (o : os) s
      | o == Multiplication = checkEquation' ns os (s * n)
      | o == Addition = checkEquation' ns os (s + n)
      | o == Concatenation = checkEquation' ns os (read (show s ++ show n) :: Int)
      | otherwise = error "Unknown operator!"

solveWith :: [Operator] -> Equation -> Maybe [[Operator]]
solveWith ops (Equation (left, right)) = 
  let solvedOps = filter (checkEquation (Equation (left, right))) (replicateM (length right - 1) ops)
   in case solvedOps of
        [] -> Nothing
        s -> Just s

solve :: Equation -> Maybe [[Operator]]
solve = solveWith [Multiplication, Addition]

part1 :: String -> IO ()
part1 s =
  let eqs = (fst . fromJust . parse equationList) s
      solvedEquations = map (\(Equation (left, _), _) -> left) (filter snd (map (\eq -> (eq, (isJust . solve) eq)) eqs `using` parList rpar))
   in print $ sum solvedEquations
  where
    equationList :: Parser [Equation]
    equationList = many1 (const <$> equation <*> char '\n')

-- | Part 2
solveConcat :: Equation -> Maybe [[Operator]]
solveConcat = solveWith [Multiplication, Addition, Concatenation]

part2 :: String -> IO ()
part2 s =
  let eqs = (fst . fromJust . parse equationList) s
      solvedEquations = map (\(Equation (left, _), _) -> left) (filter snd (map (\eq -> (eq, (isJust . solveConcat) eq)) eqs `using` parList rpar))
   in print $ sum solvedEquations
  where
    equationList :: Parser [Equation]
    equationList = many1 (const <$> equation <*> char '\n')

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