import Data.Char
import Data.List
import Data.List.Split

type Id = Int
type Len = Int
data Block = Data Id Len | Free Len deriving (Show, Eq)
type HDD = [Block]

getDataFree :: HDD -> Id -> [Len] -> HDD
getDataFree acc id (d:[]) =
  acc ++ [Data id d]
getDataFree acc id (d:f:r) =
  getDataFree (acc ++ [Data id d, Free f]) (id + 1) r

consume :: String -> HDD
consume file =
  getDataFree [] 0 $ map digitToInt $ trim file
  where trim = dropWhileEnd isSpace . dropWhile isSpace

frag1 :: HDD -> HDD -> HDD -> HDD
frag1 acc dt (Free _:rest) =
  frag1 acc dt rest
frag1 acc (Data di dl:dr) p@(Data pi pl:_)
  | di == pi = (Data pi pl:acc)
  | otherwise = frag1 (Data di dl:acc) dr p
frag1 acc (Free f:rest) (Data id d:pack)
  | d < f = frag1 (Data id d:acc) (Free (f - d):rest) pack
  | d == f = frag1 (Data id d:acc) rest pack
  | d > f = frag1 (Data id f:acc) rest (Data id (d - f):pack)

frag :: HDD -> HDD
frag hdd = reverse $ frag1 [] hdd (reverse hdd)

replace :: Block -> HDD -> HDD
replace d@(Data _ len) hdd =
  let (fore:aft:[]) = splitOn [d] hdd in
    fore ++ [Free len] ++ aft

moveUp :: HDD -> HDD -> Block -> HDD
moveUp acc (h@(Data hid _):rest) d@(Data did _)
  | hid == did = (reverse acc) ++ (h:rest)
  | otherwise = moveUp (h:acc) rest d
moveUp acc (f@(Free fl):rest) d@(Data id dl)
  | fl == dl = remove (d:acc) rest d
  | fl > dl = remove (Free (fl - dl):d:acc) rest d
  | otherwise = moveUp (f:acc) rest d
  where
    remove acc rest d = (reverse acc) ++ (replace d rest)

defrag1 :: HDD -> HDD -> HDD
defrag1 hdd []            = hdd
defrag1 hdd (Free _:rest) = defrag1 hdd rest
defrag1 hdd (d:rest)      = defrag1 (moveUp [] hdd d) rest

defrag :: HDD -> HDD
defrag hdd = defrag1 hdd (reverse hdd)

checksum1 :: Int -> Int -> HDD -> Int
checksum1 acc _ [] =
  acc
checksum1 acc ix (Free len:rest) =
  checksum1 acc (ix + len) rest
checksum1 acc ix (Data id len:rest) =
  checksum1 (acc + (sum $ map (id*) [ix..ix+len-1])) (ix + len) rest

checksum :: HDD -> Int
checksum = checksum1 0 0

part1 :: HDD -> Int
part1 = checksum . frag

part2 :: HDD -> Int
part2 = checksum . defrag

main = do
  file <- readFile "day09.input"
  let hdd = consume file
  putStrLn ("Part 1: " ++ (show $ part1 hdd))
  putStrLn ("Part 2: " ++ (show $ part2 hdd))