#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [attoparsec])"
#! nix-shell -i "ghcid -c 'ghci -Wall' -T main"

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative
import Control.Monad
import Data.Attoparsec.Text (Parser, parseOnly)
import Data.Char (isSpace)
import Data.Map.Strict (Map)
import Data.Maybe (catMaybes)
import qualified Data.Attoparsec.Text as A
import qualified Data.Map.Strict as M
import qualified Data.Text as T

-- The expected fields are as follows:
--
--   - byr (Birth Year)
--   - iyr (Issue Year)
--   - eyr (Expiration Year)
--   - hgt (Height)
--   - hcl (Hair Color)
--   - ecl (Eye Color)
--   - pid (Passport ID)
--   - cid (Country ID)

-- Count the number of valid passports - those that have all required fields.
-- Treat cid as optional. In your batch file, how many passports are valid?

data Field = BYR | IYR | EYR | HGT | HCL | ECL | PID | CID
  deriving (Show, Ord, Eq)

data PassportSpec = PassportSpec { birthYear      :: !String
                                 , issueYear      :: !String
                                 , expirationYear :: !String
                                 , height         :: !String
                                 , hairColor      :: !String
                                 , eyeColor       :: !String
                                 , passportID     :: !String
                                 , countryID      :: !(Maybe String)
                                 }
  deriving Show

fieldParser :: Parser (Map Field String)
fieldParser = ("byr:" *> A.takeTill isSpace >>= pure . (M.singleton BYR) . T.unpack)
          <|> ("iyr:" *> A.takeTill isSpace >>= pure . (M.singleton IYR) . T.unpack)
          <|> ("eyr:" *> A.takeTill isSpace >>= pure . (M.singleton EYR) . T.unpack)
          <|> ("hgt:" *> A.takeTill isSpace >>= pure . (M.singleton HGT) . T.unpack)
          <|> ("hcl:" *> A.takeTill isSpace >>= pure . (M.singleton HCL) . T.unpack)
          <|> ("ecl:" *> A.takeTill isSpace >>= pure . (M.singleton ECL) . T.unpack)
          <|> ("pid:" *> A.takeTill isSpace >>= pure . (M.singleton PID) . T.unpack)
          <|> ("cid:" *> A.takeTill isSpace >>= pure . (M.singleton CID) . T.unpack)

fieldsParser :: Parser (Map Field String)
fieldsParser = mconcat <$> fieldParser `A.sepBy` (A.space)

batchFileParser :: Parser [Map Field String]
batchFileParser = fieldsParser `A.sepBy` (A.string "\n\n")

parseFields :: String -> Either String [Map Field String]
parseFields str = parseOnly batchFileParser (T.pack str)

mkPassport1 :: Map Field String -> Maybe PassportSpec
mkPassport1 m = do
  byr <- M.lookup BYR m
  iyr <- M.lookup IYR m
  eyr <- M.lookup EYR m
  hgt <- M.lookup HGT m
  hcl <- M.lookup HCL m
  ecl <- M.lookup ECL m
  pid <- M.lookup PID m
  cid <- Just $ M.lookup CID m
  pure $ PassportSpec byr iyr eyr hgt hcl ecl pid cid

testData :: [String]
testData = [ "ecl:gry pid:860033327 eyr:2020 hcl:#fffffd"
           , "byr:1937 iyr:2017 cid:147 hgt:183cm"
           , ""
           , "iyr:2013 ecl:amb cid:350 eyr:2023 pid:028048884"
           , "hcl:#cfa07d byr:1929"
           , ""
           , "hcl:#ae17e1 iyr:2013"
           , "eyr:2024"
           , "ecl:brn pid:760753108 byr:1931"
           , "hgt:179cm"
           , ""
           , "hcl:#cfa07d eyr:2025 pid:166559648"
           , "iyr:2011 ecl:brn hgt:59in"
           ]

testData2 :: [String]
testData2 =
  [ "eyr:1972 cid:100"
  , "hcl:#18171d ecl:amb hgt:170 pid:186cm iyr:2018 byr:1926"
  , ""
  , "iyr:2019"
  , "hcl:#602927 eyr:1967 hgt:170cm"
  , "ecl:grn pid:012533040 byr:1946"
  , ""
  , "hcl:dab227 iyr:2012"
  , "ecl:brn hgt:182cm pid:021572410 eyr:2020 byr:1992 cid:277"
  , ""
  , "hgt:59cm ecl:zzz"
  , "eyr:2038 hcl:74454a iyr:2023"
  , "pid:3556412378 byr:2007"
  , ""
  , "pid:087499704 hgt:74in ecl:grn iyr:2012 eyr:2030 byr:1980"
  , "hcl:#623a2f"
  , ""
  , "eyr:2029 ecl:blu cid:129 byr:1989"
  , "iyr:2014 pid:896056539 hcl:#a97842 hgt:165cm"
  , ""
  , "hcl:#888785"
  , "hgt:164cm byr:2001 iyr:2015 cid:88"
  , "pid:545766238 ecl:hzl"
  , "eyr:2022"
  , ""
  , "iyr:2010 hgt:158cm hcl:#b6652a ecl:blu byr:1944 eyr:2021 pid:093154719"
  ]

solveDay4Part1 :: String -> Either String Int
solveDay4Part1 s =
  length . catMaybes . map mkPassport1 <$> parseFields s

eitherToMaybe :: Either a b -> Maybe b
eitherToMaybe (Right x) = Just x
eitherToMaybe (Left _) = Nothing

nothingIfFalse :: Bool -> Maybe Bool
nothingIfFalse False = Nothing
nothingIfFalse True = Just True

mkPassport2 :: Map Field String -> Maybe PassportSpec
mkPassport2 m = do
  byr <- M.lookup BYR m
  iyr <- M.lookup IYR m
  eyr <- M.lookup EYR m
  hgt <- M.lookup HGT m
  hcl <- M.lookup HCL m
  ecl <- M.lookup ECL m
  pid <- M.lookup PID m
  cid <- Just $ M.lookup CID m
  _ <- nothingIfFalse (checkByr byr)
  _ <- nothingIfFalse (checkIyr iyr)
  _ <- nothingIfFalse (checkEyr eyr)
  _ <- (checkHgt hgt)
  _ <- (checkHcl hcl)
  _ <- (checkEcl ecl)
  _ <- (checkPid pid)
  pure $ PassportSpec byr iyr eyr hgt hcl ecl pid cid
  where
    checkByr s = (\n -> n >= 1920 && n <= 2002) $ (read s :: Int)
    checkIyr s = (\n -> n >= 2010 && n <= 2020) $ (read s :: Int)
    checkEyr s = (\n -> n >= 2020 && n <= 2030) $ (read s :: Int)
    checkHcl = eitherToMaybe . parseOnly (
      A.char '#' *> (A.hexadecimal :: Parser Int) *> pure True) . T.pack
    checkEcl = eitherToMaybe . parseOnly (
      (   A.string "amb"
      <|> A.string "blu"
      <|> A.string "brn"
      <|> A.string "gry"
      <|> A.string "grn"
      <|> A.string "hzl"
      <|> A.string "oth"
      ) *> pure True) . T.pack
    checkPid = eitherToMaybe . parseOnly (A.count 9 A.digit *> A.endOfInput *> pure True) . T.pack

checkHgt :: String -> Maybe Bool
checkHgt = eitherToMaybe . parseOnly hgtParser . T.pack

hgtParser :: Parser Bool
hgtParser = (decimalIn <* A.string "in") <|> (decimalCm <* A.string "cm")

decimalIn :: Parser Bool
decimalIn = do
  d <- A.decimal
  unless (d >= (59::Int) && d <= (76::Int)) (fail "")
  pure $ True

decimalCm :: Parser Bool
decimalCm = do
  d <- A.decimal
  unless (d >= (150::Int) && d <= (193::Int)) (fail "")
  pure $ True

solveDay4Part2 :: String -> Either String Int
solveDay4Part2 s =
  length . catMaybes . map mkPassport2 <$> parseFields s

main :: IO ()
main = do
  putStrLn "Day 4 - Part 1"
  print $ unlines testData
  print $ solveDay4Part1 (unlines testData)
  input <- readFile "day4/input"
  print $ solveDay4Part1 input
  print $ solveDay4Part2 (unlines testData2)
  print $ solveDay4Part2 input