module Main where

import Data.Bifunctor (first)
import Data.Char
import Data.List
import Data.Foldable
import Data.Maybe
import Data.Functor.Identity
import Control.Monad.State

import Text.ParserCombinators.ReadP hiding (get)

import Aoc

main :: IO ()
main = do
  content <- lines <$> input 10

  print $ part1 content
  putStr $ part2 content
  where
    state = catMaybes . flip evalState (0, 1) . sequence . map instruction . parse
    part1 = sum . calculateSignalStrength . state
    part2 = unlines . splitCRTLines . drawCRT . ((1,1):) . map (first (+1)) . state

calculateSignalStrength :: [(Cycle, RegX)] -> [Int]
calculateSignalStrength = go 20
  where
    go :: Int -> [(Cycle, RegX)] -> [Int]
    go c@220 xs = singleton . strength c $ every20th c xs
    go cycle xs = (strength cycle $ every20th cycle xs) : (go (cycle + 40) . snd $ every20th cycle xs)
    strength cycle = (*cycle) . snd . last . fst
    every20th cycle = span ((<cycle) . fst)

data Instruction = AddX Int | Noop

drawCRT :: [(Cycle, RegX)] -> String
drawCRT [] = mempty
drawCRT ((c, x):[]) = foldr' (\a b -> (drawPixel a x) : b) [] [c..240]
drawCRT ((ac, ax):r@(bc, bx):xs) = (foldr' (\a b -> (drawPixel a ax) : b) [] [ac..bc-1]) ++ drawCRT (r:xs)

splitCRTLines :: String -> [String]
splitCRTLines = go
  where
    go :: String -> [String]
    go [] = [[]]
    go x = (fst $ crtLine x) : go (snd $ crtLine x)
    crtLine = splitAt 40

drawPixel :: Cycle -> RegX -> Char
drawPixel x y
  | index <= 1 = '#'
  | otherwise = '.'
  where
    index = abs (column - y)
    column = (x - 1) `mod` 40

cycles :: Instruction -> Int
cycles (AddX _) = 2
cycles Noop = 1

type Cycle = Int
type RegX = Int

instruction :: Instruction -> StateT (Cycle, RegX) Identity (Maybe (Cycle, RegX))
instruction x = do
  (cycle, regx) <- get
  let newCycle = cycle + cycles x
  case x of
    Noop -> do
      put (newCycle, regx)
      pure Nothing
    (AddX a) -> do
      let newSt = (newCycle, regx + a)
      put newSt
      pure $ Just newSt

parseNoop :: ReadP Instruction
parseNoop = do
  string "noop"
  pure Noop

parseAddX :: ReadP Instruction
parseAddX = do
  string "addx"
  skipSpaces
  n <- read <$> munch1 num
  pure $ AddX n
  where
    num c = isDigit c || c == '-'

parseInstruction :: ReadP Instruction
parseInstruction = choice [parseNoop, parseAddX]

parse :: [String] -> [Instruction]
parse = map (fst . head . readP_to_S parseInstruction)