{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import           Control.Applicative       (liftA3)
import           Control.Exception         (try)
import           Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import           Data.Bool                 (bool)
import           Data.Foldable             (fold)
import           Data.Functor              (($>))
import           Data.List                 (intersperse)
import           Data.Maybe                (maybe)
import           Data.Monoid               ((<>))
import           Data.Text                 (null, pack, strip, unpack)
import           Data.Time.Clock.POSIX     (getPOSIXTime)
import           Prelude                   hiding (FilePath, putStrLn)
import           Text.Printf               (printf)
import           Data.Text.IO              (putStrLn)
import           Turtle                    hiding (empty, fold, printf)
import           Data.Maybe
import           System.Environment (getArgs)

data Report = Report {
    uploadRate   :: UploadRate,
    downloadRate :: DownloadRate
  } deriving Show

newtype DownloadRate = DownloadRate { runDownloadRate :: TransferRate } deriving Show
newtype UploadRate = UploadRate { runUploadRate :: TransferRate } deriving Show

data TransferRate = TransferRate {
    value :: Double,
    unit  :: TransferUnit
  } deriving (Show, Eq)

data TransferUnit = Bs | KBs | MBs deriving (Show, Eq, Enum, Ord)

data Record = Record {
    timestamp     :: Timestamp,
    bytesReceived :: TotalBytesIn,
    bytesSent     :: TotalBytesOut
  } deriving Show

newtype TotalBytesIn = TotalBytesIn { runTotalBytesIn :: Integer } deriving (Show, Read)
newtype TotalBytesOut = TotalBytesOut { runTotalBytesOut :: Integer } deriving (Show, Read)
type Timestamp = Integer
type NetworkInterface = Text

dataPathForInterface :: NetworkInterface -> FilePath
dataPathForInterface interface = dataPath </> fromText interface

dataPath :: FilePath
dataPath = "/dev/shm/i3blocks-bandwidth-module"

main :: IO ()
main = do
  icons <- getArgs
  sh
          $   defaultInterface
          >>= maybe handleNoInterface (runScript icons)
          >>= liftIO
          .   putStrLn
  where
    handleNoInterface = return "No UP interface"

initPath :: NetworkInterface -> Shell FilePath
initPath interface = do
        mktree dataPath
        touch $ dataPathForInterface interface
        return $ dataPathForInterface interface

runScript :: [String] -> NetworkInterface -> Shell Text
runScript icons interface = do
  _ <- initPath interface
  record <- readRecord interface
  maybe handleNoRecordAvailable handleRecord record
    where
      handleRecord oldRecord =
          formatReport icons
          .   applyBestUnit
          .   speedReport oldRecord
          <$> writeRecord interface
      handleNoRecordAvailable = do
        _ <- writeRecord interface
        return $ "No data for " <> (pack $ show interface)

writeRecord :: NetworkInterface -> Shell Record
writeRecord interface =
        liftA3 Record
               (liftIO poxisTimeAsInteger)
               (readBytesReceived interface)
               (readBytesTransfered interface)
                >>= liftIO
                .   writeRecordToFile
    where
        recordToText record =
                pack
                        . fold
                        . intersperse " "
                        $ [ show . timestamp $ record
                          , show . runTotalBytesIn . bytesReceived $ record
                          , show . runTotalBytesOut . bytesSent $ record
                          , "\n"
                          ]
        poxisTimeAsInteger = round <$> getPOSIXTime
        writeRecordToFile record =
                ( writeTextFile (dataPathForInterface interface)
                        . recordToText
                        $ record
                        )
                        $> record

readRecord :: NetworkInterface -> Shell (Maybe Record)
readRecord interface = liftIO . runMaybeT $ extractRecord =<< safeReadTextFile
    where
        safeReadTextFile =
                MaybeT
                        $   tryToMaybe
                        <$> try
                                    (strip <$> readTextFile
                                            (dataPathForInterface interface)
                                    )
        extractRecord =
                MaybeT . return . matchToMaybe . match (decimal `sepBy` " ")
        matchToMaybe = \case
                [[time', received, sent]] -> Just $ Record
                        time'
                        (TotalBytesIn received)
                        (TotalBytesOut sent)
                _ -> Nothing
        tryToMaybe :: Either IOError b -> Maybe b
        tryToMaybe (Left  _)  = Nothing
        tryToMaybe (Right x') = Just x'

speedReport :: Record -> Record -> Report
speedReport oldRecord newRecord = Report (UploadRate uploadRate')
                                         (DownloadRate downloadRate')
    where
        timediff = max (timestamp newRecord - timestamp oldRecord) 1
        bytesReceivedDiff =
                runTotalBytesIn (bytesReceived newRecord)
                        - runTotalBytesIn (bytesReceived oldRecord)
        bytesSentDiff =
                runTotalBytesOut (bytesSent newRecord)
                        - runTotalBytesOut (bytesSent oldRecord)
        downloadRate' = TransferRate
                (fromIntegral bytesReceivedDiff / fromIntegral timediff)
                Bs
        uploadRate' = TransferRate
                (fromIntegral bytesSentDiff / fromIntegral timediff)
                Bs

getUploadIcon :: [String] -> Maybe Text
getUploadIcon icons = case icons of
  uploadIcon : _ -> Just $ pack uploadIcon
  _              -> Nothing

getDownloadIcon :: [String] -> Maybe Text
getDownloadIcon icons = case icons of
  _ : downloadIcon : _ -> Just $ pack downloadIcon
  _                    -> Nothing

formatReport :: [String] -> Report -> Text
formatReport icons report =
  let downloadIcon = fromMaybe "Rx " $ getDownloadIcon icons
      uploadIcon   = fromMaybe "Tx " $ getUploadIcon icons
  in  downloadIcon
                <> ( pack
                   . printf "%.1f"
                   . value
                   . runDownloadRate
                   . downloadRate
                   $ report
                   )
                <> (pack . show . unit . runDownloadRate . downloadRate $ report)
                <> "  "
                <> uploadIcon
                <> ( pack
                   . printf "%.1f"
                   . value
                   . runUploadRate
                   . uploadRate
                   $ report
                   )
                <> (pack . show . unit . runUploadRate . uploadRate $ report)

applyBestUnit :: Report -> Report
applyBestUnit = liftA2
        Report
        (UploadRate . applyBestUnit' . runUploadRate . uploadRate)
        (DownloadRate . applyBestUnit' . runDownloadRate . downloadRate)
    where
        applyBestUnit' transferRate | value transferRate < 1024 = transferRate
        applyBestUnit' transferRate = applyBestUnit'
                $ convertRate transferRate (succ . unit $ transferRate)

convertRate :: TransferRate -> TransferUnit -> TransferRate
convertRate rate to | unit rate < to = convertRate (convertRateUp rate) to
    where
        convertRateUp (TransferRate rate' from)
          = TransferRate (rate' / 1024) (succ from)
convertRate rate' to | unit rate' > to = convertRate (convertRateDown rate') to
    where
        convertRateDown (TransferRate rate from)
          = TransferRate (rate * 1024) (pred from)
convertRate rate _ = rate

defaultInterface :: Shell (Maybe NetworkInterface)
defaultInterface =
        textToMaybe . strip
                <$> strict (inshell
                            "ip addr | awk '/state UP/ {print $2}' | sed 's/.$//'"
                            mempty
                    )
        where textToMaybe text' = bool (Just text') Nothing (Data.Text.null text')

readBytesTransfered :: Text -> Shell TotalBytesOut
readBytesTransfered interface =
        TotalBytesOut . read . unpack . strip
                <$> strict (inshell
                            (  "cat /sys/class/net/"
                            <> interface
                            <> "/statistics/tx_bytes"
                            )
                            mempty
                    )

readBytesReceived :: Text -> Shell TotalBytesIn
readBytesReceived interface =
        TotalBytesIn . read . unpack . strip
                <$> strict (inshell
                            (  "cat /sys/class/net/"
                            <> interface
                            <> "/statistics/rx_bytes"
                            )
                            mempty
                    )