plexuploader

plex-upload.hs

#!/usr/bin/env stack
{- stack
    --resolver lts-12.14
    --install-ghc
    runghc
    --package HSH
    --package cmdargs
    --package pretty-simple
-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RecordWildCards       #-}

import           Control.Monad
import           Data.Char
import           Data.Function
import           Data.List
import           Data.Time.Clock
import           Debug.Trace
import           HSH
import           System.Console.CmdArgs.Implicit as A
import           System.Directory
import           System.FilePath
import           Text.Pretty.Simple

data Paths = Paths
  { localPath    :: Maybe FilePath
  , remotePath   :: FilePath
  , overlayPath  :: FilePath
  , pruneSeasons :: Bool
  , rcloneDest   :: String
  , remoteRoot   :: String
  , maxTime      :: Int
  , minSeasonAge :: Int
  , minOldAge    :: Int
  , oldFiles     :: Bool
  } deriving (Show, Data, Typeable)

defPaths =
  Paths
    { localPath =
        def &= typDir &=
        help "Location where the local copies of files are kept" &=
        explicit &=
        A.name "local-path" &=
        A.name "l"
    , remotePath = def &= typ "RemotePath" &= argPos 1
    , overlayPath = def &= typ "OverlayPath" &= argPos 0
    , pruneSeasons =
        False &= help "Move all but the most recent season to remote storage" &=
        explicit &=
        A.name "prune-seasons" &=
        A.name "p"
    , rcloneDest =
        "plex" &= help "rclone destination name" &= explicit &=
        A.name "rclone-dest" &=
        A.name "d"
    , remoteRoot =
        "/" &= help "rclone destination root folder" &= explicit &=
        A.name "rclone-root" &=
        A.name "r"
    , maxTime =
        360 &= help "Maximum time (in miniutes) to spend uploading files"
    , minSeasonAge =
        30 &= help "Minimum age to be pruned by the season pruner, in days"
    , minOldAge =
        60 &= help "Minimum age to be pruned by the old files pruner, in days"
    , oldFiles = False &= help "Do not remove old Files"
    } &=
  program "upload.hs" &=
  summary "PlexUploader v0.0.0, (C) Nathan McCarty 2018" &=
  details
    [ "Moves older, less likely to be watched video files from local to cloud storage"
    , ""
    , "RemotePath specifies the location where the remote/cloud storage is mounted. OverlayPath specifies the location where the Overlay/Union of the remote and local storage is located."
    , ""
    ] &=
  verbosity

data Config = Config
  { localPath    :: FilePath
  , remotePath   :: FilePath
  , overlayPath  :: FilePath
  , pruneSeasons :: Bool
  , rcloneRoot   :: String
  , maxTime      :: Int
  , minSeasonAge :: Int
  , minOldAge    :: Int
  , oldFiles     :: Bool
  } deriving (Show)

main = do
  paths <- cmdArgs defPaths
  let config = pathsToConfig paths
  let Config {..} = config
  whenLoud $ putStrLn "Working Configuration:"
  whenLoud $ pPrint config
  setCurrentDirectory overlayPath
  titles <- makeTitleList overlayPath localPath
  titleStats titles
  final <- selectFiles config titles
  let limit = fromInteger . toInteger $ maxTime * 60 :: NominalDiffTime
  putStrLn $ "Time Limit is " ++ (show limit)
  tally <- mapUntil limit (uploadEpisode config) final
  let numFiles = length tally
  putStrLn $ "\nUploaded " ++ (show numFiles) ++ " files."

--
-- Program IO Blocks
--

titleStats :: [Title] -> IO ()
titleStats titles = do
  putStrLn $ "Analyzing " ++ (show . length $ titles) ++ " series."
  putStrLn $ "  - " ++ (show . sum . map titleSeasons $ titles) ++ " seasons."
  putStrLn $ "  - " ++ (show . sum . map titleEpisodes $ titles) ++ " episodes."
  putStrLn $
    "  - " ++ (show . sum . map titleOnDisk $ titles) ++ " episodes on disk.\n"

selectFiles :: Config -> [Title] -> IO [Episode]
selectFiles Config {..} titles = do
  let stage1 =
        if pruneSeasons
          then filter ((> minSeasonAge) . fileAge) . weave . map oldSeasons $
               titles
          else []
  putStrLn $
    (show . length $ stage1) ++ " Episodes selected for season pruning."
  let stage2 =
        if oldFiles
          then oldEpisodes titles minOldAge
          else []
  putStrLn $
    (show . length $ stage2) ++ " Episodes selected for old age pruning."
  let combined = nub . weave $ [stage1, stage2]
  putStrLn $ (show . length $ combined) ++ " Unique episodes to be uploaded.\n"
  return combined

timeIO :: IO a -> IO (a, NominalDiffTime)
timeIO action = do
  start <- getCurrentTime
  a <- action
  end <- getCurrentTime
  return (a, diffUTCTime end start)

mapUntil :: NominalDiffTime -> (a -> IO b) -> [a] -> IO [b]
mapUntil limit function items =
  if limit <= 0 || null items
    then return []
    else do
      let first:rest = items
      (head, time) <- timeIO . function $ first
      let newLimit = limit - time
      putStrLn $ "Time Remaining is " ++ (show newLimit)
      tail <- mapUntil newLimit function rest
      return $ head : tail

--
-- File Upload Mechanics
--

uploadEpisode :: Config -> Episode -> IO ()
uploadEpisode Config {..} episode = do
  let path = (location :: Episode -> FilePath) episode
  let local = localPath </> path
  let remote = rcloneRoot ++ (dropFileName path)
  putStrLn $ "Uploading " ++ path
  run (("rclone", ["move",local,remote]) :: (String,[String]))

--
-- File Selectors
--

-- | Merges any number of lists together
weave :: [[a]] -> [a]
weave [] = []
weave lists =
  let nlists = filter (not . null) lists
      front = map head nlists
      back = weave . map tail $ nlists
   in front ++ back

-- | List of episodes to be removed based on being part of a non-current season,
-- oldest episodes first. Will only provide files that are on disk.
oldSeasons :: Title -> [Episode]
oldSeasons title =
  let s = seasons title
      sorted = sortBy (compare `on` number) s
      olds = init sorted
   in filter isOnDisk . concat . map episodes $ olds

-- | Returns the number of days old a date is
daysOld :: UTCTime -> UTCTime -> Int
daysOld currentTime time =
  let diff = diffUTCTime currentTime time
   in floor $ diff / nominalDay

-- | Returns a list of episods older than a specificed age
-- sorted by age in descending order
oldEpisodes :: [Title] -> Int -> [Episode]
oldEpisodes titles minAge =
  let eps = concat . map episodes . concat . map seasons $ titles
      sortedEps = sortBy (flip (compare `on` fileAge)) eps
   in filter isOnDisk . takeWhile ((> minAge) . fileAge) $ eps

--
-- House Keeping
--

-- The stuff in this section is primarly for dealing with the command line arguments and getting the program ready to run

-- | Turns a paths object into a config object.
-- Will set localPath = overlayPath if not set
pathsToConfig :: Paths -> Config
pathsToConfig Paths {localPath = Just lp, ..} =
  Config
    { localPath = lp
    , remotePath
    , overlayPath
    , pruneSeasons
    , rcloneRoot = (rcloneDest ++ ":" ++ remoteRoot)
    , maxTime
    , minSeasonAge
    , minOldAge
    , oldFiles = not oldFiles
    }
pathsToConfig Paths {..} =
  Config
    { localPath = overlayPath
    , remotePath
    , overlayPath
    , pruneSeasons
    , rcloneRoot = (rcloneDest ++ ":" ++ remoteRoot)
    , maxTime
    , minSeasonAge
    , minOldAge
    , oldFiles = not oldFiles
    }

--
-- Data Structures
--

data DiskState
  = All
  | Some
  | None
  deriving (Show, Eq)

combineStates :: DiskState -> DiskState -> DiskState
combineStates All All   = All
combineStates All _     = Some
combineStates Some _    = Some
combineStates None None = None
combineStates None _    = Some

compositeState :: [DiskState] -> DiskState
compositeState = foldl combineStates All

isAll :: DiskState -> Bool
isAll All = True
isAll _   = False

data Episode = Episode
  { name     :: FilePath
  , location :: FilePath
  , onDisk   :: DiskState
  , modTime  :: UTCTime
  , fileAge  :: Int
  } deriving (Show, Eq)

isOnDisk :: Episode -> Bool
isOnDisk = isAll . (onDisk :: Episode -> DiskState)

data Season = Season
  { number   :: Int
  , name     :: FilePath
  , location :: FilePath
  , episodes :: [Episode]
  , onDisk   :: DiskState
  }

seasonEpisodes :: Season -> Int
seasonEpisodes = length . episodes

seasonOnDisk :: Season -> Int
seasonOnDisk = length . filter isOnDisk . episodes

data Title = Title
  { name     :: FilePath
  , location :: FilePath
  , seasons  :: [Season]
  , onDisk   :: DiskState
  }

titleSeasons :: Title -> Int
titleSeasons = length . seasons

titleEpisodes :: Title -> Int
titleEpisodes = sum . map seasonEpisodes . seasons

titleOnDisk :: Title -> Int
titleOnDisk = sum . map seasonOnDisk . seasons

--
-- Scraping the directory tree and producing a [Title]
--

-- | Helper function for makeSeason and makeTitle
ioBuilder :: (a -> b -> IO c) -> IO a -> b -> IO c
ioBuilder f ia b = do
  a <- ia
  d <- f a b
  return d

-- | Produces an Episode given a location, a season, and a title
-- Will check the localFiles to see if it is on disk. Expects a relative path
makeEpisode :: FilePath -> FilePath -> IO Episode
makeEpisode localDir location = do
  let name = takeFileName location
  let localLocation = localDir </> location
  exists <- doesFileExist localLocation
  let diskState =
        if exists
          then All
          else None
  currentTime <- getCurrentTime
  modTime <- getModificationTime location
  let tmpFileAge = daysOld currentTime modTime
  let fileAge =
        if tmpFileAge < 0
          then 1000
          else tmpFileAge -- If the file was made in the future, the usenet downloader messed up. Lets pretend its very old
  return Episode {name, location, onDisk = diskState, modTime, fileAge}

-- | Will produce a season when given a season directory(name) and a title
-- Expects a relative to libary root path.
-- Requires both overlay and local file directories to be provided
makeSeason :: FilePath -> FilePath -> FilePath -> IO Season
makeSeason overlayDir localDir location = do
  let name = takeFileName location
  let number = read . filter isDigit $ name
  episodeFiles <- listDirectory (overlayDir </> location)
  let episodeLocations = map (location </>) episodeFiles
  episodes <- mapM (makeEpisode localDir) episodeLocations
  let od = compositeState $ map (onDisk :: Episode -> DiskState) episodes
  return Season {number, name, location, episodes, onDisk = od}

-- | Will Produce a title when given a title directory
-- Expects a relative to library root path.
-- Requires both overlay and local file directories to be provided
makeTitle :: FilePath -> FilePath -> FilePath -> IO Title
makeTitle overlayDir localDir location = do
  let name = takeFileName location
  seasonFiles <- listDirectory (overlayDir </> location)
  let seasonLocations = map (location </>) seasonFiles
  seasons <- mapM (makeSeason overlayDir localDir) seasonLocations
  let od = compositeState . map (onDisk :: Season -> DiskState) $ seasons
  return Title {name, location, seasons, onDisk = od}

-- | Will Produce a list of titles given an overlay directory
-- also requires the local directory for sane on disk detection
makeTitleList :: FilePath -> FilePath -> IO [Title]
makeTitleList overlayDir localDir = do
  titleLocations <- listDirectory overlayDir
  mapM (makeTitle overlayDir localDir) titleLocations