module Day08 where
import Data.Char (ord)
import Data.List (sort, permutations)
import Data.List.Split (splitOn)
import Data.List.Ordered (isSorted)
import Data.Maybe (mapMaybe, isJust, fromJust)

toMaybe :: a -> Bool -> Maybe a
toMaybe x True = Just x
toMaybe x False = Nothing

type Input = ([String], [String])
countToSeg :: Int -> Maybe Int
countToSeg 2 = Just 1
countToSeg 4 = Just 4
countToSeg 3 = Just 7
countToSeg 7 = Just 8
countToSeg _ = Nothing

signalsToDigit :: [Int] -> Maybe Int
signalsToDigit [0,1,2,4,5,6] = Just 0
signalsToDigit [2,5] = Just 1
signalsToDigit [0,2,3,4,6] = Just 2
signalsToDigit [0,2,3,5,6] = Just 3
signalsToDigit [1,2,3,5] = Just 4
signalsToDigit [0,1,3,5,6] = Just 5
signalsToDigit [0,1,3,4,5,6] = Just 6
signalsToDigit [0,2,5] = Just 7
signalsToDigit [0,1,2,3,4,5,6] = Just 8
signalsToDigit [0,1,2,3,5,6] = Just 9
signalsToDigit x
    | not $ isSorted x = signalsToDigit $ sort x
    | otherwise = Nothing


segToNum :: Char -> Int
segToNum x = ord x - ord 'a'

parseLine :: String -> Input
parseLine s = (words signal, words output)
    where
        [signal, output] = splitOn "|" s

countEasy :: Input -> Int
countEasy (signal,output) = length $ mapMaybe (countToSeg . length) output

testDigit :: [Int] -> String -> Bool
testDigit perm digit = isJust $ signalsToDigit $ map (\x -> perm !! segToNum x) digit

testPermutation :: [Int] -> [String] -> Bool
testPermutation perm = all (testDigit perm)

selectPermutation :: [[Int]] -> [String] -> [Int]
selectPermutation perms digits = head $ filter(`testPermutation` digits) perms

decode :: [[Int]] -> Input -> Int
decode p (signals,output) = read $ concatMap show digits
    where
        digits = map (fromJust . (signalsToDigit . map (\x -> perm !! segToNum x))) output
        perm = selectPermutation p signals

main :: IO ()
main = do
    input <- map parseLine . lines <$> getContents

    putStrLn $ "part 1: " ++ show (sum $ map countEasy input)

    let perms = permutations [0..6]
    let valid_perms = map (decode perms) input
    -- mapM_ print valid_perms
    putStrLn $ "part 2: " ++ show (sum valid_perms)