module Main where

import System.Environment (getArgs)
import Control.Monad (void)
import Data.List (transpose)

import Text.Parsec (parse, Parsec, endBy, many, char, anyChar, (<|>), try, string, lookAhead, getPosition, sourceName)
import Text.Parsec.Char (oneOf)
import Text.ParserCombinators.Parsec (eof)

rows :: Parsec String () [String]
rows = endBy (many $ oneOf "XMAS") (void $ char '\n') <* eof

rotate90 :: [[a]] -> [[a]]
rotate90 = reverse . transpose

diagonals :: [[a]] -> [[a]]
diagonals = (++)
    -- Lower half (excluding the main diagonal)
    <$> reverse . transpose . map reverse . zipWith take [1..] . drop 1
    -- Upper half
    <*> transpose . zipWith drop [0..]

-- | counts the number of "XMAS" (as well as the backwards "SAMX") in the input
countLineXMAS :: Parsec String () Int
countLineXMAS =
    try (lookAhead (string "XMAS" <|> string "SAMX")
        >> anyChar >> (+1) <$> countLineXMAS)
    <|> (anyChar >> countLineXMAS)
    <|> (eof >> return 0)

countAllXMAS :: Parsec String () Int
countAllXMAS = do
    grid <- rows
    let linesOfInterest = grid ++ concatMap ($ grid) [transpose, diagonals, diagonals . rotate90]
    sName <- sourceName <$> getPosition
    counts <- mapM (either (\p -> fail $ "Parse Error! " ++ show p) return . parse countLineXMAS sName) linesOfInterest
    return $ sum counts

part1 :: FilePath -> IO Int
part1 fp = do
    content <- readFile fp
    case parse countAllXMAS fp content of
        Left e -> fail $ "Oh no! " ++ show e
        Right v -> return v

main :: IO ()
main = do
    args <- getArgs
    (fp:_) <- return args
    resPart1 <- part1 fp
    putStrLn $ "Solution (Part 1) : " ++ show resPart1