shieeeeet.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

-- {-# LANGUAGE TDNR #-}
module Main where

import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Crypto.Hash.MD5 as MD5
import Data.ByteString
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy as LBS
import Data.Char
import Data.Hex
import Data.List as List
import Data.Maybe
import Data.Monoid
import Data.Set as Set
import Network.HTTP.Types
import Network.Wai as Wai
import Network.Wai.Handler.Warp
import Network.Wreq
import System.Directory
import System.Process
import Control.Monad.Trans.Except
import Control.Lens
import Control.Arrow ((>>>))

main :: IO ()
main = do
  ops <- newTVarIO Set.empty
  run 3000 $ app ops

type OpId = String

app :: TVar (Set OpId) -> Request -> (Wai.Response -> IO b) -> IO b
app ops req respond = respond =<< serveTranscode ops req

badParam :: ByteString -> Wai.Response
badParam name = responseLBS status400 [] $ LBS.fromStrict $ "bad " <> name

serveTranscode :: TVar (Set OpId) -> Request -> IO Wai.Response
serveTranscode ops req = do
  e <- runExceptT $ do
    i <- queryValue "i"
    f <- queryValue "f"
    let outputName = C.unpack $ getOutputName i opt f
    return $ bracket_ (claimOp outputName ops) (releaseOp outputName ops) $ do
      ready <- doesFileExist outputName
      unless ready $ transcode outputName i opt
      return $ responseFile status200 [] outputName Nothing
  mergeEither e
  where
    qs = queryString req
    queryValue :: ByteString -> ExceptT (IO Wai.Response) IO ByteString
    queryValue k = case getFirstQueryValue k qs of
      Just v -> return v
      Nothing -> throwE $ return $ badParam k
    opt = getQueryValues "opt" qs

maybeToEither l Nothing = Left l
maybeToEither _ (Just r) = Right r

mergeEither :: Either a a -> a
mergeEither e = case e of
  Left v -> v
  Right v -> v

transcode name i opts = do
  callProcess (List.head args) (List.tail args)
  where
    args = ffmpegArgs name (C.unpack i) $ List.map C.unpack opts

download i file = do
  r <- get i
  LBS.writeFile file $ r ^. responseBody

ffmpegArgs outputName i opts =
  ["nice", "ffmpeg", "-hide_banner", "-i", i] ++ opts ++ ["-y", outputName]

getFirstQueryValue :: ByteString -> Query -> Maybe ByteString
getFirstQueryValue key = List.find (\(k, _) -> k == key) >>> fmap snd >>> join

getQueryValues :: ByteString -> Query -> [ByteString]
getQueryValues key = mapMaybe f
  where
    f (k, Just v) =
      if k == key
        then Just v
        else Nothing
    f (_, Nothing) = Nothing

claimOp :: String -> TVar (Set OpId) -> IO ()
claimOp op ops =
  atomically $ do
    opsval <- readTVar ops
    if member op opsval
      then retry
      else modifyTVar ops $ Set.insert op

releaseOp :: String -> TVar (Set OpId) -> IO ()
releaseOp file active = atomically $ modifyTVar active $ Set.delete file

getOutputName :: ByteString -> [ByteString] -> ByteString -> ByteString
getOutputName i opts f =
  (C.pack . List.map toLower . C.unpack . hex $ hashStrings (i : opts)) <> "." <>
  f

hashStrings :: [ByteString] -> ByteString
hashStrings = updates MD5.init >>> finalize