import           Data.List.Split (splitOn)
import           Data.Vector     (Vector, fromList, (!), (//))

operations = [(+), (*)]

parseOpcode :: Int -> [Int]
parseOpcode 0    = []
parseOpcode code = [mod code 10] ++ parseOpcode (code `div` 10)

executeMemory :: Vector Int -> Int -> IO (Maybe (Vector Int))
executeMemory memory ip = do
  case (opcode !! 0) of
    1 -> addOrMul
    2 -> addOrMul
    3 -> input -- Would like to use `in` here, but thats a reserved keyword in Haskell
    4 -> out
    5 -> jnz
    6 -> jz
    7 -> cmpLT
    8 -> eq
    9 ->
      if ((length opcode > 1) && (opcode !! 1) == 9)
        then return $ Just memory
        else return Nothing
    _ -> return Nothing
  where
    opcode = parseOpcode $ memory ! ip
    op = operations !! ((opcode !! 0) - 1)
    addOrMul =
      let addr1 = memory ! (ip + 1)
          addr2 = memory ! (ip + 2)
          addr3 = memory ! (ip + 3)
          param1 =
            if (length opcode > 2 && opcode !! 2 == 1)
              then addr1
              else memory ! addr1
          param2 =
            if (length opcode > 3 && opcode !! 3 == 1)
              then addr2
              else memory ! addr2
          result = op param1 param2
       in executeMemory (memory // [(addr3, result)]) (ip + 4)
    input =
      let addr1 = memory ! (ip + 1)
       in do putStr "INPUT>> "
             input <- getLine
             executeMemory
               (memory // [(addr1, (read :: String -> Int) input)])
               (ip + 2)
    out =
      let addr1 = memory ! (ip + 1)
       in do putStrLn $
               (++) "OUTPUT>> " $
               show $
               if (length opcode > 2 && opcode !! 2 == 1)
                 then addr1
                 else memory ! addr1
             executeMemory memory (ip + 2)
    jnz =
      let addr1 = memory ! (ip + 1) -- Zero flag
          addr2 = memory ! (ip + 2) -- New IP position
          param1 =
            if (length opcode > 2 && opcode !! 2 == 1)
              then addr1
              else memory ! addr1
          param2 =
            if (length opcode > 3 && opcode !! 3 == 1)
              then addr2
              else memory ! addr2
          newIP =
            if (param1 /= 0)
              then param2
              else (ip + 3)
       in executeMemory memory newIP
    jz =
      let addr1 = memory ! (ip + 1) -- Zero flag
          addr2 = memory ! (ip + 2) -- New IP position
          param1 =
            if (length opcode > 2 && opcode !! 2 == 1)
              then addr1
              else memory ! addr1
          param2 =
            if (length opcode > 3 && opcode !! 3 == 1)
              then addr2
              else memory ! addr2
          newIP =
            if (param1 == 0)
              then param2
              else (ip + 3)
       in executeMemory memory newIP
    cmpLT =
      let addr1 = memory ! (ip + 1)
          addr2 = memory ! (ip + 2)
          addr3 = memory ! (ip + 3)
          param1 =
            if (length opcode > 2 && opcode !! 2 == 1)
              then addr1
              else memory ! addr1
          param2 =
            if (length opcode > 3 && opcode !! 3 == 1)
              then addr2
              else memory ! addr2
          modMem =
            if (param1 < param2)
              then memory // [(addr3, 1)]
              else memory // [(addr3, 0)]
       in executeMemory modMem (ip + 4)
    eq =
      let addr1 = memory ! (ip + 1)
          addr2 = memory ! (ip + 2)
          addr3 = memory ! (ip + 3)
          param1 =
            if (length opcode > 2 && opcode !! 2 == 1)
              then addr1
              else memory ! addr1
          param2 =
            if (length opcode > 3 && opcode !! 3 == 1)
              then addr2
              else memory ! addr2
          modMem =
            if (param1 == param2)
              then memory // [(addr3, 1)]
              else memory // [(addr3, 0)]
       in executeMemory modMem (ip + 4)

main :: IO ()
main = do
  input <-
    map (read :: String -> Int) <$> splitOn "," <$> readFile "day5input.txt"
  let memory = fromList input
  result <- executeMemory memory 0
  putStrLn "DUMP MEMORY (Y/N)?"
  putStr "INPUT>> "
  dumpMem <- getLine
  case dumpMem of
    "Y" -> putStrLn $ show result
    _   -> return ()