module Main where

import Data.Char (digitToInt)
import Data.Either (rights)
import Data.List (foldl')
import Data.Map (Map)
import qualified Data.Map as M
import System.Environment (getArgs)

-- | A candidate is either a left-most number we've just selected, or
-- a complete candidate represented as a pair.
type Candidate = Either Int (Int, Int)

main :: IO ()
main = do
  input <- readFile . head =<< getArgs
  putStrLn $ "Part 1: " ++ show (part1 input)
  putStrLn $ "Part 2: " ++ show (part2 input)

-- | Maintain a list of candidates (either partial, or complete) and
-- then select the greatest one for each line.
--
-- Insight: By keeping a list of possible candidates we can always
-- backtrack out of an accidental corner (like making a partial
-- candidate of the last digit).
part1 :: String -> Int
part1 = sum . map (maximum . map combine . rights . foldl go [] . map digitToInt) . lines
  where
    go :: [Candidate] -> Int -> [Candidate]
    go [] x = pure $ Left x
    go (Left a : candidates') x
      | x > a = Left x : Right (a, x) : candidates'
      | otherwise = Right (a, x) : candidates'
    go candidates@(Right (a, b) : candidates') x
      | x > a = Left x : Right (a, x) : candidates
      | x > b = Right (a, x) : candidates'
      | otherwise = candidates

    combine :: (Int, Int) -> Int
    combine (a, b) = a * 10 + b

type CandidateN = [Int]

-- | Maintain one partial candidates for each length.
--
-- Insight: If there are at least N-1 digits left after we've spotted
-- the greatest digit, that digit needs to be the first digit of our
-- candidates.
--
-- Insight 2: We only ever need to store one candidate per length.
part2 :: String -> Int
part2 = sum . map (combine . (M.! 12) . foldl' checkPotential mempty . map digitToInt) . lines
  where
    checkPotential :: Map Int CandidateN -> Int -> Map Int CandidateN
    checkPotential cs x =
      let partials = filter ((< 12) . length) $ M.elems cs
          partials' = [x] : map (genPartial x) partials
       in foldl' insertMaxCandidate cs $ filter ((<= 12) . length) partials'

    -- generates new partials by appending x at candidate or after
    -- replacing a shedding smaller digits.
    genPartial :: Int -> CandidateN -> CandidateN
    genPartial x candidate = candidate ++ [x]

    -- replaces a candidate with a better one with the same length
    insertMaxCandidate :: Map Int CandidateN -> CandidateN -> Map Int CandidateN
    insertMaxCandidate m c = M.insertWith maxCandidate (length c) c m

    maxCandidate :: CandidateN -> CandidateN -> CandidateN
    maxCandidate a b
      | combine a > combine b = a
      | otherwise = b

    combine :: CandidateN -> Int
    combine xs = sum $ zipWith power [0 :: Int ..] (reverse xs)

    power :: Int -> Int -> Int
    power pow b = b * 10 ^ pow

testInput :: String
testInput =
  unlines
    [ "987654321111111",
      "811111111111119",
      "234234234234278",
      "818181911112111"
    ]