module Main (main) where

import Data.List (findIndex, groupBy)
import Data.Maybe (fromJust)

-- | Part 1
data DiskItem = File Int Int | Free Int deriving (Show, Eq)

isFree :: DiskItem -> Bool
isFree (Free _) = True
isFree _ = False

isFile :: DiskItem -> Bool
isFile (File _ _) = True
isFile _ = False

fileLen :: DiskItem -> Int
fileLen (File _ len) = len
fileLen (Free len) = len

fileId :: DiskItem -> Int
fileId (File fId _) = fId
fileId (Free _) = error "Free has no id!"

readDiskMap :: String -> [DiskItem]
readDiskMap = readDiskMap' 0
  where
    readDiskMap' :: Int -> String -> [DiskItem]
    readDiskMap' _ [] = []
    readDiskMap' _ ['\n'] = []
    readDiskMap' fileId (i : is)
      | even fileId = File (fileId `div` 2) (read [i] :: Int) : readDiskMap' (fileId + 1) is
      | otherwise = Free (read [i] :: Int) : readDiskMap' (fileId + 1) is

expandDiskMap :: [DiskItem] -> [DiskItem]
expandDiskMap [] = []
expandDiskMap ((File fileId len) : is) = replicate len (File fileId 1) ++ expandDiskMap is
expandDiskMap ((Free len) : is) = replicate len (Free 1) ++ expandDiskMap is

showDiskMap :: [DiskItem] -> IO ()
showDiskMap [] = putStrLn ""
showDiskMap ((File fileId len) : is) = putStr (concat [show fileId | _ <- [1 .. len]]) >> showDiskMap is
showDiskMap ((Free len) : is) = putStr (replicate len '.') >> showDiskMap is

defragmentStep :: [DiskItem] -> [DiskItem]
defragmentStep diskMap =
  let diskSize = length diskMap
      firstFree = fromJust $ findIndex isFree diskMap
      lastFile = diskSize - 1 - fromJust (findIndex (not . isFree) (reverse diskMap))
   in take firstFree diskMap ++ [diskMap !! lastFile] ++ take (lastFile - firstFree - 1) (drop (firstFree + 1) diskMap) ++ [Free 1] ++ take (diskSize - lastFile) (drop (lastFile + 1) diskMap)

defragmentWith :: ([DiskItem] -> [DiskItem]) -> [DiskItem] -> [DiskItem]
defragmentWith f diskMap
  | 2 == length (groupBy (\a b -> isFile a && isFile b || isFree a && isFree b) diskMap) = diskMap
  | otherwise = defragment (f diskMap)

defragment :: [DiskItem] -> [DiskItem]
defragment = defragmentWith defragmentStep

checksum :: [DiskItem] -> Int
checksum = checksum' 0
  where
    checksum' :: Int -> [DiskItem] -> Int
    checksum' _ [] = 0
    checksum' c ((Free len) : rest) = checksum' (c + len) rest
    checksum' c ((File fId len) : rest) = fId * sum (take len [c..]) + checksum' (c + len) rest

part1 :: String -> IO ()
part1 s =
  let diskMap = (expandDiskMap . readDiskMap) s
   in print $ checksum $ defragment diskMap

-- | Part 2
freeCandidates :: [DiskItem] -> [(DiskItem, [DiskItem])]
freeCandidates [] = []
freeCandidates ((File f len) : ds) = (File f len, []) : freeCandidates ds
freeCandidates ((Free len) : ds) = (Free len, filter ((&&) <$> ((len >=) . fileLen) <*> isFile) ds) : freeCandidates ds

replaceItem :: Int -> [DiskItem] -> [DiskItem]
replaceItem _ [] = []
replaceItem fId (Free n : ds) = Free n : replaceItem fId ds
replaceItem fId (File f len : ds)
  | fId == f = Free len : replaceItem fId ds
  | otherwise = File f len : replaceItem fId ds

defragmentStep2 :: [DiskItem] -> [DiskItem]
defragmentStep2 diskMap =
  let candidates = freeCandidates diskMap
      replacePos = findIndex (not . null . snd) candidates
   in case replacePos of
        Nothing -> diskMap
        Just pos ->
          let replacement = (last . snd) (candidates !! pos)
              freeLen = (fileLen . fst) (candidates !! pos)
           in take pos diskMap
                ++ [replacement]
                ++ ([Free (freeLen - fileLen replacement) | (freeLen - fileLen replacement) /= 0])
                ++ take (length diskMap - pos) (drop (pos + 1) (replaceItem (fileId replacement) diskMap))

defragment2 :: [DiskItem] -> [DiskItem]
defragment2 diskMap =
  let stepped = defragmentStep2 diskMap
  in if stepped == diskMap then diskMap else defragment2 stepped

part2 :: String -> IO ()
part2 s =
  let diskMap = readDiskMap s
   in print $ checksum $ defragment2 diskMap

-- | Main
main :: IO ()
main = readFile "resources/day9.txt" >>= part2