-- https://adventofcode.com/2023/day/1

import           Data.Bifunctor               (bimap, second)
import           Data.Maybe                   (fromJust, fromMaybe, isJust)
import           Data.Text                    (Text, pack, replace, unpack)
import           Text.ParserCombinators.ReadP (ReadP, choice, many, readP_to_S,
                                               string, (<++))

main :: IO ()
main = do
  entries <- lines <$> readFile "src/advent_of_code/2023/input1"
  print "Advent of Code 2023 - Day 1"
  -- print entries

  print ("Part 1: " <> show (solveP1 entries))
  print ("Part 2: " <> show (solveP2 entries))

--
-- Part 1
--

solveP1 :: [String] -> Int
solveP1 = sum . map solveP1Line

solveP1Line :: String -> Int
solveP1Line = sumPairS . bimap head last . doubleString . filter isNumP1 . tokenized tokensP1

tokensP1 :: ReadP String
tokensP1 = parseNumberChars <++ parseChars

--
-- Part 2
--

solveP2 :: [String] -> Int
solveP2 = sum . map solveP2Line

solveP2Line :: String -> Int
solveP2Line = sumPairS . bimap head last . doubleString . map stringToNum . filter isNumP2 . tokenized tokensP2 . replaceAll overlaps

tokensP2 :: ReadP String
tokensP2 = parseNumberWords <++ parseNumberChars <++ parseChars

--
-- Utility functions
--

-- Doing this because I want to use a bifunctor
doubleString :: str -> (str, str)
doubleString    str =  (str, str)

sumPairS :: (String, String) -> Int
sumPairS (x,y) = read (x ++ y)

stringToNum :: String -> String
stringToNum str = fromMaybe str (lookup str $ zip numberWords numberChars)

--
-- Boilerplate equality checks
-- Now condensed
--

isNumP1 :: String -> Bool
isNumP1 = flip elem numberChars

isNumP2 :: String -> Bool
isNumP2 = flip elem (numberChars ++ numberWords)

--
-- Parsing for question
--

letterChars :: String
letterChars = ['a'..'z']

numberChars :: [String]
numberChars = map show [0..9]

numberWords :: [String]
numberWords = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]

parseChars :: ReadP String
parseChars = choice $ map (\c -> string [c]) letterChars

parseNumberChars :: ReadP String
parseNumberChars = choice $ map string numberChars

parseNumberWords :: ReadP String
parseNumberWords = choice $ map string numberWords

tokenized :: ReadP a -> String -> [a]
tokenized tokenSet = fst . last . readP_to_S (many tokenSet)

--
-- Detecting overlaps
--   - Instead of reversing to tokenize from the right, I'm using the numberWords list to create a set of possible string overlaps
--   -   later I'll detect these and replace them, to be able to parse all strings from left to right
--

overlaps :: [(String, String)]
overlaps = concatMap combineOverlaps overlapsDerived

-- Not used, but the full list written manually
overlapsSpec :: [((String, String), String)]
overlapsSpec =
  [ (("zero" , "one"  ), "01")
  , (("one"  , "eight"), "18")
  , (("two"  , "one"  ), "21")
  , (("three", "eight"), "38")
  , (("five" , "eight"), "58")
  , (("seven", "nine" ), "79")
  , (("eight", "two"  ), "82")
  , (("eight", "three"), "83")
  , (("nine" , "eight"), "98")
  ]

overlapsDerived :: [(String, [String])]
overlapsDerived = map (second fromJust) . filter (isJust . snd) . zip numberWords $ map findOverlaps numberWords
  where
    findOverlaps :: String -> Maybe [String]
    findOverlaps str = if null overlaps then Nothing else Just overlaps
      where
        overlaps = filter (\s -> last str == head s) numberWords

combineOverlaps :: (String, [String]) -> [(String, String)]
combineOverlaps (numStr, []       ) = []
combineOverlaps (numStr, over:laps) = (numStr ++ tail over, stringToNum numStr ++ stringToNum over) : combineOverlaps (numStr, laps)

-- String replace using Text lol
replace' :: String -> String -> String -> String
replace' needle replacement haystack = unpack $ replace needle' replacement' haystack'
  where
    needle'      = pack needle
    replacement' = pack replacement
    haystack'    = pack haystack

replaceAll :: [(String, String)] -> String -> String
replaceAll overlaps str = foldl (\haystack (needle, replacement) -> replace' needle replacement haystack) str overlaps

--
-- Examples
--

example1 :: [String]
example1 =
  [ "1abc2"
  , "pqr3stu8vwx"
  , "a1b2c3d4e5f"
  , "treb7uchet"
  ]

example2 :: [String]
example2 =
  [ "two1nine"
  , "eightwothree"
  , "abcone2threexyz"
  , "xtwone3four"
  , "4nineeightseven2"
  , "zoneight234"
  , "7pqrstsixteen"
  ]