-- | Copyright: (c) 2021 The closed eye of love
-- SPDX-License-Identifier: BSD-3-Clause
-- Maintainer: Poscat <poscat@mail.poscat.moe>, berberman <berberman@yandex.com>
-- Stability: alpha
-- Portability: portable
-- A set of utilities for downloading pixiv things.
module Web.Pixiv.Download
  ( -- * DownloadM monad
    DownloadM,
    liftMaybe,
    liftToPixivT,
    runDownloadM,

    -- * Download actions
    downloadPixiv,
    downloadSingleIllust,
    downloadUgoiraToMP4,
  )
where

import Control.Lens
import Control.Monad.Reader
import Control.Monad.Trans.Maybe
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import Network.HTTP.Client
import System.FilePath ((</>))
import System.IO.Temp (getCanonicalTemporaryDirectory, withTempDirectory)
import System.Process
import Web.Pixiv.Types
import Web.Pixiv.Types.Lens
import Web.Pixiv.Types.PixivT
import Web.Pixiv.Utils

-- | 'DownloadM' monad is a synonym for 'IO' computation wrapped in 'Manager' environment,
-- which may exit without producing value, indicating the download failed.
type DownloadM = MaybeT (ReaderT Manager IO)

-- | Lifts a pure 'Maybe' value to 'DownloadM'.
liftMaybe :: Maybe a -> DownloadM a
liftMaybe :: Maybe a -> DownloadM a
liftMaybe Maybe a
x = ReaderT Manager IO (Maybe a) -> DownloadM a
forall (m :: Type -> Type) a. m (Maybe a) -> MaybeT m a
MaybeT (ReaderT Manager IO (Maybe a) -> DownloadM a)
-> (IO (Maybe a) -> ReaderT Manager IO (Maybe a))
-> IO (Maybe a)
-> DownloadM a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Maybe a) -> ReaderT Manager IO (Maybe a)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> DownloadM a) -> IO (Maybe a) -> DownloadM a
forall a b. (a -> b) -> a -> b
$ Maybe a -> IO (Maybe a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Maybe a
x

-- | Executes the download computation in 'IO'.
runDownloadM :: Manager -> DownloadM a -> IO (Maybe a)
runDownloadM :: Manager -> DownloadM a -> IO (Maybe a)
runDownloadM Manager
manager DownloadM a
m = DownloadM a -> ReaderT Manager IO (Maybe a)
forall (m :: Type -> Type) a. MaybeT m a -> m (Maybe a)
runMaybeT DownloadM a
m ReaderT Manager IO (Maybe a)
-> (ReaderT Manager IO (Maybe a) -> IO (Maybe a)) -> IO (Maybe a)
forall a b. a -> (a -> b) -> b
& (ReaderT Manager IO (Maybe a) -> Manager -> IO (Maybe a))
-> Manager -> ReaderT Manager IO (Maybe a) -> IO (Maybe a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT Manager IO (Maybe a) -> Manager -> IO (Maybe a)
forall r (m :: Type -> Type) a. ReaderT r m a -> r -> m a
runReaderT Manager
manager

-- | Lifts a download computation to 'PixivT'.
--
-- 'DownloadM' needs 'Manager' to perform the download, which can be provided by 'TokenState'.
liftToPixivT :: (MonadIO m) => DownloadM a -> PixivT m (Maybe a)
liftToPixivT :: DownloadM a -> PixivT m (Maybe a)
liftToPixivT DownloadM a
m = do
  PixivState {tokenState :: PixivState -> TokenState
tokenState = TokenState {UTCTime
Manager
Token
manager :: TokenState -> Manager
expirationTime :: TokenState -> UTCTime
refreshToken :: TokenState -> Token
accessToken :: TokenState -> Token
manager :: Manager
expirationTime :: UTCTime
refreshToken :: Token
accessToken :: Token
..}} <- PixivT m PixivState
forall (m :: Type -> Type). MonadPixiv m => m PixivState
readPixivState
  IO (Maybe a) -> PixivT m (Maybe a)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe a) -> PixivT m (Maybe a))
-> IO (Maybe a) -> PixivT m (Maybe a)
forall a b. (a -> b) -> a -> b
$ Manager -> DownloadM a -> IO (Maybe a)
forall a. Manager -> DownloadM a -> IO (Maybe a)
runDownloadM Manager
manager DownloadM a
m

-- | Downloads something in 'DownloadM', given url.
downloadPixiv :: Text -> DownloadM LBS.ByteString
downloadPixiv :: Text -> DownloadM ByteString
downloadPixiv Text
url = do
  let addReferer :: Request -> Request
addReferer Request
r = Request
r {requestHeaders :: RequestHeaders
requestHeaders = [(HeaderName
"Referer", ByteString
"https://app-api.pixiv.net/")]}
  Manager
manager <- MaybeT (ReaderT Manager IO) Manager
forall r (m :: Type -> Type). MonadReader r m => m r
ask
  Request
req <- Request -> Request
addReferer (Request -> Request)
-> MaybeT (ReaderT Manager IO) Request
-> MaybeT (ReaderT Manager IO) Request
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (String -> MaybeT (ReaderT Manager IO) Request
forall (m :: Type -> Type). MonadThrow m => String -> m Request
parseRequest (String -> MaybeT (ReaderT Manager IO) Request)
-> (Text -> String) -> Text -> MaybeT (ReaderT Manager IO) Request
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack (Text -> MaybeT (ReaderT Manager IO) Request)
-> Text -> MaybeT (ReaderT Manager IO) Request
forall a b. (a -> b) -> a -> b
$ Text
url)
  Response ByteString
resp <- IO (Response ByteString)
-> MaybeT (ReaderT Manager IO) (Response ByteString)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (Response ByteString)
 -> MaybeT (ReaderT Manager IO) (Response ByteString))
-> IO (Response ByteString)
-> MaybeT (ReaderT Manager IO) (Response ByteString)
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ByteString)
httpLbs Request
req Manager
manager
  ByteString -> DownloadM ByteString
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (ByteString -> DownloadM ByteString)
-> ByteString -> DownloadM ByteString
forall a b. (a -> b) -> a -> b
$ Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
resp

-- | Downloads a single page illust.
--
-- Chooses the first one if the illust has many pages,
-- preferring high quality images. Returns `Nothing` if can't find any image url.
downloadSingleIllust :: Illust -> DownloadM LBS.ByteString
downloadSingleIllust :: Illust -> DownloadM ByteString
downloadSingleIllust Illust
i = do
  Text
url <- Maybe Text -> DownloadM Text
forall a. Maybe a -> DownloadM a
liftMaybe (Maybe Text -> DownloadM Text) -> Maybe Text -> DownloadM Text
forall a b. (a -> b) -> a -> b
$ Illust -> [Text]
extractImageUrlsFromIllust Illust
i [Text] -> Getting (First Text) [Text] Text -> Maybe Text
forall s a. s -> Getting (First a) s a -> Maybe a
^? Getting (First Text) [Text] Text
forall s a. Cons s s a a => Traversal' s a
_head
  Text -> DownloadM ByteString
downloadPixiv Text
url

-- | Downloads 'UgoiraFrame's, then converts it to MP4 calling external @ffmpeg@.
downloadUgoiraToMP4 ::
  -- | Information of ugoira to download
  UgoiraMetadata ->
  -- | Path to @ffmpeg@
  Maybe FilePath ->
  DownloadM (String, LBS.ByteString)
downloadUgoiraToMP4 :: UgoiraMetadata -> Maybe String -> DownloadM (String, ByteString)
downloadUgoiraToMP4 UgoiraMetadata
meta (String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"ffmpeg" -> String
ffmpeg) = do
  let ffconcat :: ByteString
ffconcat = UgoiraMetadata -> ByteString
ugoiraMetadataToFFConcat UgoiraMetadata
meta
  ByteString
bs0 <- Text -> DownloadM ByteString
downloadPixiv (Text -> DownloadM ByteString) -> Text -> DownloadM ByteString
forall a b. (a -> b) -> a -> b
$ UgoiraMetadata
meta UgoiraMetadata -> Getting Text UgoiraMetadata Text -> Text
forall s a. s -> Getting a s a -> a
^. (ZipUrls -> Const Text ZipUrls)
-> UgoiraMetadata -> Const Text UgoiraMetadata
forall s a. HasZipUrls s a => Lens' s a
zipUrls ((ZipUrls -> Const Text ZipUrls)
 -> UgoiraMetadata -> Const Text UgoiraMetadata)
-> ((Text -> Const Text Text) -> ZipUrls -> Const Text ZipUrls)
-> Getting Text UgoiraMetadata Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Const Text Text) -> ZipUrls -> Const Text ZipUrls
forall s a. HasZipMedium s a => Lens' s a
zipMedium
  String
systmp <- IO String -> MaybeT (ReaderT Manager IO) String
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO IO String
getCanonicalTemporaryDirectory
  IO (String, ByteString) -> DownloadM (String, ByteString)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (String, ByteString) -> DownloadM (String, ByteString))
-> IO (String, ByteString) -> DownloadM (String, ByteString)
forall a b. (a -> b) -> a -> b
$
    String
-> String
-> (String -> IO (String, ByteString))
-> IO (String, ByteString)
forall (m :: Type -> Type) a.
(MonadMask m, MonadIO m) =>
String -> String -> (String -> m a) -> m a
withTempDirectory String
systmp String
"ugoira" ((String -> IO (String, ByteString)) -> IO (String, ByteString))
-> (String -> IO (String, ByteString)) -> IO (String, ByteString)
forall a b. (a -> b) -> a -> b
$ \String
temp -> do
      String -> ByteString -> IO ()
unzipArchive String
temp ByteString
bs0
      let concatFilePath :: String
concatFilePath = String
temp String -> String -> String
</> String
"concat"
          convertProcess :: CreateProcess
convertProcess =
            ( String -> [String] -> CreateProcess
proc
                String
ffmpeg
                [ String
"-y",
                  String
"-i",
                  String
"concat",
                  String
"-c:v",
                  String
"libx264",
                  String
"-vf",
                  String
"pad=ceil(iw/2)*2:ceil(ih/2)*2",
                  String
"-pix_fmt",
                  String
"yuv420p",
                  String
"-lossless",
                  String
"1",
                  String
"ugoira.mp4"
                ]
            )
              { cwd :: Maybe String
cwd = String -> Maybe String
forall a. a -> Maybe a
Just String
temp
              }
      String -> ByteString -> IO ()
BS.writeFile String
concatFilePath ByteString
ffconcat
      (ExitCode
_code, String
_stdout, String
stderr) <- CreateProcess -> String -> IO (ExitCode, String, String)
readCreateProcessWithExitCode CreateProcess
convertProcess String
""
      (String
stderr,) (ByteString -> (String, ByteString))
-> IO ByteString -> IO (String, ByteString)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO ByteString
LBS.readFile (String
temp String -> String -> String
</> String
"ugoira.mp4")