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

windows3 :: [a] -> [(a, a, a)]
windows3 (a:b:c:tl) = (a, b, c) : windows3 (b:c:tl)
windows3 _ = []

windows3x3 :: [[a]] -> [((a, a, a), (a, a, a), (a, a, a))]
windows3x3 ls = [(wa, wb, wc)
    | (a, b, c) <- windows3 ls
    , (wa, wb, wc) <- zip3 (windows3 a) (windows3 b) (windows3 c)]

countCrossMAS :: [String] -> Int
countCrossMAS = sum
    . map (\((tl, _, tr),
             (_ , c, _ ),
             (bl, _, br)) ->
        if isMAS [tl, c, br] && isMAS [tr, c, bl] then 1 else 0)
    . windows3x3
    where
        isMAS :: String -> Bool
        isMAS = (||) <$> (==) "MAS" <*> (==) "SAM"

part2 :: FilePath -> IO Int
part2 fp = do
    content <- readFile fp
    case parse rows fp content of
        Left e -> fail $ "Oh no! " ++ show e
        Right r -> return $ countCrossMAS r

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