content: {-# LANGUAGE LambdaCase #-}
  
  import           Control.Concurrent          (forkIO, setNumCapabilities)
  import           Control.Concurrent.Chan     (Chan, newChan, readChan,
                                                writeChan, writeList2Chan)
  import           Control.Exception           (IOException, catch)
  import           Control.Monad               (filterM, forever, join,
                                                replicateM, replicateM_, (>=>))
  import           Control.Parallel.Strategies (parTraversable, rseq, using,
                                                withStrategy)
  import           Data.Char                   (isDigit)
  import           Data.List                   (find, sort)
  import qualified Data.Map.Strict             as Map
  import           GHC.Conc                    (getNumProcessors)
  import           System.Directory            (doesDirectoryExist, doesFileExist,
                                                listDirectory,
                                                pathIsSymbolicLink)
  import           System.Environment          (getArgs)
  import           System.FilePath.Posix       ((</>))
  import           System.IO                   (FilePath, IOMode (ReadMode),
                                                hFileSize, hPutStrLn, stderr,
                                                withFile)
  import           Text.Printf                 (hPrintf, printf)
  
  data Item = File FilePath Integer | Folder FilePath deriving (Show)
  
  type FGKey = (Integer, Integer)
  type FrequencyGroup = (FGKey, Integer)
  type FrequencyGroups = Map.Map FGKey Integer
  
  newFrequencyGroups :: FrequencyGroups
  newFrequencyGroups = Map.empty
  
  fileSizes :: [Item] -> [Integer]
  fileSizes = foldr f [] where f (File _ n) acc = n:acc
                               f _          acc = acc
  
  folders :: [Item] -> [FilePath]
  folders = foldr f [] where f (Folder p) acc = p:acc
                             f _          acc = acc
  
  totalBytes :: [Item] -> Integer
  totalBytes = sum . fileSizes
  
  counts :: [Item] -> (Integer, Integer)
  counts = foldr (\x (a, b) -> case x of File _ _ -> (succ a, b)
                                         Folder _ -> (a, succ b)) (0, 0)
  
  -- |Creates 'FrequencyGroups' from the provided size and data set.
  frequencyGroups :: Int             -- ^ Desired number of frequency groups.
                  -> [Integer]       -- ^ List of collected file sizes. Must be sorted.
                  -> FrequencyGroups -- ^ Returns a 'FrequencyGroups' for the file sizes.
  frequencyGroups _ [] = newFrequencyGroups
  frequencyGroups totalGroups xs
    | length xs == 1 = Map.singleton (head xs, head xs) 1
    | otherwise = foldr placeGroups newFrequencyGroups xs `using` parTraversable rseq
    where
      range = maximum xs - minimum xs
      groupSize = succ $ ceiling $ realToFrac range / realToFrac totalGroups
      groups = takeWhile (<=groupSize + maximum xs) $ iterate (+groupSize) 0
      groupMinMax = zip groups (pred <$> tail groups)
      findGroup n = find (\(low, high) -> n >= low && n <= high)
  
      incrementCount (Just n) = Just (succ n) -- Update count for range.
      incrementCount Nothing  = Just 1        -- Insert new range with initial count.
  
      placeGroups n fgMap = case findGroup n groupMinMax of
        Just k  -> Map.alter incrementCount k fgMap
        Nothing -> fgMap -- Should never happen.
  
  expandGroups :: Int             -- ^ Desired number of frequency groups.
               -> [Integer]       -- ^ List of collected file sizes.
               -> Integer         -- ^ Computed frequency group limit.
               -> FrequencyGroups -- ^ Expanded 'FrequencyGroups'
  expandGroups gsize fileSizes groupThreshold
    | groupThreshold > 0 = loop 15 $ frequencyGroups gsize sortedFileSizes
    | otherwise = frequencyGroups gsize sortedFileSizes
    where
      sortedFileSizes = sort fileSizes
      loop 0 gs = gs -- break out in case we can't go below threshold
      loop n gs | all (<= groupThreshold) $ Map.elems gs = gs
                | otherwise = loop (pred n) (expand gs)
  
      expand :: FrequencyGroups -> FrequencyGroups
      expand = foldr f . withStrategy (parTraversable rseq) <*>
        Map.mapWithKey groupsFromGroup . Map.filter (> groupThreshold)
        where
          f :: Maybe (FGKey, FrequencyGroups) -- ^ expanded frequency group
            -> FrequencyGroups                -- ^ accumulator
            -> FrequencyGroups                -- ^ merged accumulator
          f (Just (k, fg)) acc = Map.union (Map.delete k acc) fg
          f Nothing        acc = acc
  
          groupsFromGroup
            :: FGKey                          -- ^ Group Key
            -> Integer                        -- ^ Count
            -> Maybe (FGKey, FrequencyGroups) -- ^ Returns expanded 'FrequencyGroups' with base key it replaces.
          groupsFromGroup (min, max) count
            | length range > 1 = Just ((min, max), frequencyGroups gsize range)
            | otherwise        = Nothing
            where
              range = filter (\n -> n >= min && n <= max) sortedFileSizes
  
  displaySize :: Integer -> String
  displaySize n
    |              n <= 2^10 = printf "%8dB " n
    | n >= 2^10 && n <= 2^20 = display (2^10) "KB"
    | n >= 2^20 && n <= 2^30 = display (2^20) "MB"
    | n >= 2^30 && n <= 2^40 = display (2^30) "GB"
    | n >= 2^40 && n <= 2^50 = display (2^40) "TB"
    | otherwise = "Too large!"
    where
      display :: Double -> String -> String
      display b = printf "%7.2f%s " (realToFrac n / b)
  
  displayFrequency :: Integer -> FrequencyGroup -> IO ()
  displayFrequency filesCount ((min, max), count) = do
    printf "%s <-> %s" (displaySize min) (displaySize max)
    printf "= %-10d %6.3f%%: %-5s\n" count percentage bars
    where
      percentage :: Double
      percentage = (realToFrac count / realToFrac filesCount) * 100
      size = round percentage
      bars | size == 0 = "▍"
           | otherwise = replicate size '█'
  
  folderWorker :: Chan FilePath -> Chan [Item] -> IO ()
  folderWorker folderChan resultItemsChan =
    forever (readChan folderChan >>= collectItems >>= writeChan resultItemsChan)
  
  collectItems :: FilePath -> IO [Item]
  collectItems folderPath = catch tryCollect $ \e -> do
      hPrintf stderr "Skipping: %s\n" $ show (e :: IOException)
      pure []
    where
      tryCollect = (fmap (folderPath </>) <$> listDirectory folderPath) >>=
        mapM (\p -> doesDirectoryExist p >>=
                \case True  -> pure $ Folder p
                      False -> File p <$> withFile p ReadMode hFileSize)
  
  parallelItemCollector :: FilePath -> IO [Item]
  parallelItemCollector folder = do
    wCount <- getNumProcessors
    setNumCapabilities wCount
    printf "Using %d worker threads\n" wCount
    folderChan <- newChan
    resultItemsChan <- newChan
    replicateM_ wCount (forkIO $ folderWorker folderChan resultItemsChan)
    loop folderChan resultItemsChan [Folder folder]
    where
      loop :: Chan FilePath -> Chan [Item] -> [Item] -> IO [Item]
      loop folderChan resultItemsChan xs = do
        regularFolders <- filterM (pathIsSymbolicLink >=> (pure . not)) $ folders xs
        if null regularFolders then pure []
        else do
          writeList2Chan folderChan regularFolders
          childItems <- replicateM (length regularFolders) (readChan resultItemsChan)
          result <- mapM (loop folderChan resultItemsChan) childItems
          pure (join childItems <> join result)
  
  parseArgs :: [String] -> Either String (FilePath, Int)
  parseArgs (x:y:xs)
    | all isDigit y = Right (x, read y)
    | otherwise     = Left "Invalid frequency group size"
  parseArgs (x:xs) = Right (x, 4)
  parseArgs _ = Right (".", 4)
  
  main :: IO ()
  main = parseArgs <$> getArgs >>= \case
      Left errorMessage -> hPutStrLn stderr errorMessage
      Right (path, groupSize) -> do
        items <- parallelItemCollector path
        let (fileCount, folderCount) = counts items
        printf "Total files: %d\nTotal folders: %d\n" fileCount folderCount
        printf "Total size: %s\n" $ displaySize $ totalBytes items
        printf "\nDistribution:\n\n%9s  <-> %9s %7s\n" "From" "To" "Count"
        putStrLn $ replicate 46 '-'
        let results = expandGroups groupSize (fileSizes items) (groupThreshold fileCount)
        mapM_ (displayFrequency fileCount) $ Map.assocs results
    where
      groupThreshold = round . (*0.25) . realToFrac
  