-- Copyright (c) 2015-16 Nicola Bonelli <nicola@pfq.io>
--
-- This program is free software; you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation; either version 2 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.
--
-- You should have received a copy of the GNU General Public License
-- along with this program; if not, write to the Free Software
-- Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
--

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

module Main where

import Data.Semigroup
import Data.List
import Data.List.Split
import Data.Maybe
import Data.Either
import Data.Char
import Data.Time.Clock
import qualified Data.ByteString as B (ByteString, readFile)
import qualified Data.Yaml as Y
import Data.Yaml (FromJSON(..), (.:), (.:?))

import Control.Applicative
import Control.Monad
import Control.Monad.State
import Control.Concurrent (threadDelay)
import Control.Exception

import Data.Data

import System.Console.ANSI
import System.Console.CmdArgs
import System.Directory (getHomeDirectory, doesFileExist, getModificationTime)
import System.IO.Unsafe
import System.IO.Error
import System.Process
import System.Exit
import System.FilePath
import System.Posix.Signals
import System.Posix.Types
import qualified Network.PFQ as Q (version)

proc_cpuinfo, proc_modules, tmp_pfq :: String

proc_cpuinfo = "/proc/cpuinfo"
proc_modules = "/proc/modules"
tmp_pfq      = "/tmp/pfq.opt"

dmesg, pidof, ethtool, setpci, cpufreq_set, rmmod, insmod, modprobe, lsmod, ifconfig :: String

dmesg       = "/bin/dmesg"
pidof       = "/bin/pidof"
ethtool     = "/sbin/ethtool"
setpci      = "/usr/bin/setpci"
cpufreq_set = "/usr/bin/cpufreq-set"
rmmod       = "/sbin/rmmod"
insmod      = "/sbin/insmod"
modprobe    = "/sbin/modprobe"
lsmod       = "/sbin/lsmod"
ifconfig    = "/sbin/ifconfig"


bold  = setSGRCode [SetConsoleIntensity BoldIntensity]
red   = setSGRCode [SetColor Foreground Vivid Red]
blue  = setSGRCode [SetColor Foreground Vivid Blue]
reset = setSGRCode []


configFiles = [ "/etc/pfq.conf.yaml"
              , "/root/.pfq.conf.yaml"
              , "/root/pfq.conf.yaml"
              , "/etc/pfq.conf"
              , "/root/.pfq.conf"
              , "/root/pfq.conf"
              ]

defaultDelay = 200000


newtype OptString = OptString { getOptString :: String }
    deriving (Show, Read, Eq)

newtype OptList a = OptList { getOptList :: [a] }
    deriving (Show, Read, Eq)

instance Semigroup OptString where
    a <> OptString "" = a
    _ <> b = b

instance Semigroup (OptList a) where
    a <> OptList [] = a
    _ <> b = b


data SystemError = SystemError String ExitCode
    deriving (Show, Typeable)

instance Exception SystemError


data Device =
    Device
    {   devname   :: String
    ,   devspeed  :: Maybe Int
    ,   channels  :: Maybe Int
    ,   flowctrl  :: Maybe Bool
    ,   ethopt    :: [(String, String)]
    } deriving (Show, Read, Eq)

instance FromJSON Device where
  parseJSON (Y.Object v) =
    Device <$>
        v .:   "devname"    <*>
        v .:?  "devspeed"   <*>
        v .:?  "channels"   <*>
        v .:?  "flowctrl"   <*>
        v .:   "ethopt"
  parseJSON _ = fail "Expected Object for Device value"

data Driver =
    Driver
    {   drvmod      :: String
    ,   drvopt      :: [String]
    ,   pci         :: [(String, String)]
    ,   instances   :: Int
    ,   devices     :: [Device]
    } deriving (Show, Read, Eq)


instance FromJSON Driver where
  parseJSON (Y.Object v) =
    Driver <$>
        v .:   "drvmod"    <*>
        v .:   "drvopt"    <*>
        v .:   "pci"       <*>
        v .:   "instances" <*>
        v .:   "devices"
  parseJSON _ = fail "Expected Object for Driver value"


data Config = Config
    {   pfq_module     :: String
    ,   pfq_options    :: [String]
    ,   exclude_core   :: [Int]
    ,   irq_affinity   :: [String]
    ,   cpu_governor   :: String
    ,   drivers        :: [Driver]
    } deriving (Show, Read, Eq)


instance FromJSON Config where
  parseJSON (Y.Object v) =
      Config <$>
        v .:   "pfq_module"     <*>
        v .:   "pfq_options"    <*>
        v .:   "exclude_core"   <*>
        v .:   "irq_affinity"   <*>
        v .:   "cpu_governor"   <*>
        v .:   "drivers"
  parseJSON _ = fail "Expected Object for Config value"



instance Semigroup Config where
    (Config mod1 opt1 excl1 algo1 gov1 drvs1) <> (Config mod2 opt2 excl2 algo2 gov2 drvs2) =
     Config
     {  pfq_module   = getOptString $ OptString mod1  <> OptString mod2
     ,  pfq_options  = getOptList   $ OptList opt1 <> OptList opt2
     ,  exclude_core = excl1 <> excl2
     ,  irq_affinity = algo1 <> algo2
     ,  cpu_governor = getOptString $ OptString gov1 <> OptString gov2
     ,  drivers      = drvs1 <> drvs2
     }


data Options = Options
    {   config     :: Maybe String
    ,   kmodule    :: String
    ,   algorithm  :: String
    ,   governor   :: String
    ,   first_core :: Int
    ,   exclude    :: [Int]
    ,   queues     :: Maybe Int
    ,   force      :: Bool
    ,   others     :: [String]
    } deriving (Show, Read, Data, Typeable)


options :: Mode (CmdArgs Options)
options = cmdArgsMode $ Options
    {   config     = Nothing       &= typ "FILE" &= help "Specify config file (default ~/.pfq.conf)"
    ,   kmodule    = ""            &= help "Override the kmodule specified in config file"
    ,   queues     = Nothing       &= help "Specify hardware channels (i.e. Intel RSS)"
    ,   algorithm  = ""            &= help "Irq affinity algorithm: naive, round-robin, even, odd, all-in:id, comb:id"
    ,   governor   = ""            &= help "Set cpufreq governor"
    ,   first_core = 0             &= typ "NUM" &= help "First core used for irq affinity"
    ,   exclude    = []            &= typ "CORE" &= help "Exclude core from irq affinity"
    ,   force      = False         &= help "Force PFQ reload!"
    ,   others     = []            &= args
    } &= summary ("pfq-load (PFQ " ++ Q.version ++ ")") &= program "pfq-load"


-------------------------------------------------------------------------------------------


main :: IO ()
main = handle (\(SystemError msg code) -> putStrBoldLn msg *> exitWith code) $ do

    -- load options...
    home <- getHomeDirectory
    opt  <- cmdArgsRun options
    conf <- (<> mkConfig opt) <$> loadConfig (catMaybes (config opt : map Just configFiles)) opt
    pmod <- getProcModules
    core <- getNumberOfPhyCores
    bal  <- getProcessID "irqbalance"
    frd  <- getProcessID "cpufreqd"

    -- compute the command to load pfq...

    let pfqLoadCmd = if null (pfq_module conf)
        then mkLoadModuleCmd ProbeMod  "pfq" (pfq_options conf)
        else mkLoadModuleCmd InsertMod (pfq_module conf) (pfq_options conf)


    -- determine whether pfq load is required...

    tstamp <- if null (pfq_module conf)
                then getCurrentTime
                else getModificationTime $ pfq_module conf


    pfqForceLoad <- or <$> sequence [ return $ force opt
                                    , not <$> isLoaded "pfq"
                                    , (/= (pfqLoadCmd ++ " " ++ show tstamp)) <$> possiblyReadFile tmp_pfq
                                    ]
    writeFile tmp_pfq pfqLoadCmd
    appendFile tmp_pfq (" " ++ show tstamp)

    -- check queues
    when (maybe False (> core) (queues opt)) $ errorWithoutStackTrace "queues number is too big!"

    -- unload pfq and drivers that depend on it...
    evalStateT (unloadDependencies pfqForceLoad "pfq") pmod

    -- check irqbalance deaemon
    unless (null bal) $
        putStrBoldLn ("Irqbalance daemon detected @pid " ++ show bal ++ ". Sending SIGKILL...") *> forM_ bal (signalProcess sigKILL)

    -- check cpufreqd deaemon
    unless (null frd) $
        putStrBoldLn ("Cpufreqd daemon detected @pid " ++ show frd ++ ". Sending SIGKILL...") *> forM_ frd (signalProcess sigKILL)

    -- set cpufreq governor...
    unless (null (cpu_governor conf)) $
        forM_ [0.. core-1] $ \n ->
            runSystem (cpufreq_set ++ " -g " ++ cpu_governor conf ++ " -c " ++ show n) "*** cpufreq-set error! Make sure you have cpufrequtils installed! ***"

    -- load PFQ (if required)...
    if pfqForceLoad
       then do
            putStrBoldLn (blue ++ "Loading pfq [" ++ pfqLoadCmd ++ "]")
            system $ dmesg ++ " --clear"
            runSystem pfqLoadCmd (red ++ "load pfq.ko error!")
            xs <- lines <$> readProcess dmesg ["-t"] []
            forM_ xs $ \x -> when ("pfq: unknown parameter" `isInfixOf` x) (throw $ SystemError (red ++ "load pfq.ko: " ++ x ++ "!") (ExitFailure 42))

       else putStrBoldLn $ red ++ "Using pfq.ko module already loaded (use --force to reload it)..."

    -- update current loaded proc/modules
    pmod2 <- getProcModules

    -- unload drivers...
    unless (null (drivers conf)) $
        putStrBoldLn "Unloading vanilla/standard drivers..." *>
        evalStateT (forM_ (drivers conf) $ unloadDependencies True . takeBaseName . drvmod) pmod2

    -- load and configure device drivers...
    forM_ (drivers conf) $ \drv -> do
        let rss = maybe [] (mkRssOption (drvmod drv) (instances drv)) (queues opt)
        unless (null $ drvmod drv) $
            loadModule InsertMod (drvmod drv) (drvopt drv ++ rss)
        threadDelay defaultDelay
        forM_ (devices drv) $ setupDevice (queues opt)

    -- setup per-driver pci options..
    forM_ (drivers conf) $  setupPCI . pci


    -- set interrupt affinity...
    putStrBoldLn "Setting irq affinity..."
    setupIRQAffinity (first_core opt) (exclude_core conf) (irq_affinity conf) (getDevices conf)

    putStrBoldLn "PFQ ready."


mkRssOption :: String -> Int -> Int -> [String]
mkRssOption driver ndev nqueue
    | "ixgbe.ko" `isSuffixOf` driver = [ "RSS=" ++ intercalate "," (replicate ndev (show nqueue)) ]
    | "igb.ko"   `isSuffixOf` driver = [ "RSS=" ++ intercalate "," (replicate ndev (show nqueue)) ]
    | otherwise                      = []


mkConfig :: Options -> Config
mkConfig
    Options
    {   config    = _
    ,   kmodule   = mod
    ,   algorithm = algo
    ,   exclude   = excl
    ,   governor  = gov
    ,   others    = opt } =
    Config
    {   pfq_module    = mod
    ,   pfq_options   = opt
    ,   exclude_core  = excl
    ,   irq_affinity  = [algo | not (null algo)]
    ,   cpu_governor  = gov
    ,   drivers       = []
    }


getFirstConfig :: [FilePath] -> IO (Maybe FilePath)
getFirstConfig xs = filterM doesFileExist xs >>= \case
      [ ]   -> return Nothing
      (x:_) -> return $ Just x


clean :: String -> String
clean =  unlines . filter notComment . lines
        where notComment = (not . ("#" `isPrefixOf`)) . dropWhile isSpace


loadConfig :: [FilePath] -> Options -> IO Config
loadConfig confs opt =
    getFirstConfig confs >>= \case
        Nothing   -> putStrBoldLn "Using default config..." >> return (mkConfig opt)
        Just conf ->
            putStrBoldLn ("Using " ++ conf ++ " config...") *>
            if "yaml" `isInfixOf` conf
               then do
                   conf' <- Y.decodeFileEither conf
                   case conf' of
                      Right x -> return x
                      Left  x -> throw $ SystemError (Y.prettyPrintParseException x) (ExitFailure 42)
               else read . clean <$> readFile conf


{-# NOINLINE getNumberOfPhyCores #-}
getNumberOfPhyCores :: IO Int
getNumberOfPhyCores = readFile proc_cpuinfo >>= \file ->
    return $ (length . filter (isInfixOf "processor") . lines) file


type ProcModules = [ (String, [String]) ]
type ModStateT   = StateT ProcModules


getProcModules :: IO ProcModules
getProcModules =
    readFile proc_modules >>= \file ->
        return $ map (\l -> let ts = words l in (head ts,  filter (\s -> (not . null) s && s /= "-") $
                                splitOn "," (ts !! 3))) $ lines file


rmmodFromProcMOdules :: String -> ProcModules -> ProcModules
rmmodFromProcMOdules name = filter (\(m,ds) -> m /= name )


getProcessID :: String -> IO [ProcessID]
getProcessID name = map read . words <$> catchIOError (readProcess pidof [name] "") (\_-> return [])


moduleDependencies :: String -> ProcModules -> [String]
moduleDependencies  name =
    concatMap (\(m,ds) -> if m == name then ds else [])


unloadDependencies :: Bool              -- unload the given module as well
                   -> String            -- kernel module name
                   -> ModStateT IO ()
unloadDependencies rm name = do
    proc_mods <- get
    mapM_ (unloadDependencies True) (moduleDependencies name proc_mods)
    when (rm && isModuleLoaded name proc_mods) $ do
        liftIO $ rmmod' name
        put $ rmmodFromProcMOdules name proc_mods
        where rmmod' name = putStrBoldLn ("Unloading " ++ name ++ "...") *>
                            runSystem (rmmod ++ " " ++ name) ("rmmod " ++ name ++ " error.")
              isModuleLoaded name =  any (\(mod,_) -> mod == name)


data LoadMode = InsertMod | ProbeMod
                    deriving Eq


isLoaded :: String -> IO Bool
isLoaded mod = getProcModules >>= (\pmod -> return $ mod `elem` fmap fst pmod)


-- putStrBoldLn $ "Loading " ++ name ++ " " ++ show opts


mkLoadModuleCmd :: LoadMode -> String -> [String] -> String
mkLoadModuleCmd mode name opts =
    tool ++ " " ++ name ++ " " ++ unwords opts
    where tool = if mode == InsertMod then insmod
                                      else modprobe


loadModule :: LoadMode -> String -> [String] -> IO ()
loadModule mode name opts = putStrBoldLn ("Loading " ++ name ++ " " ++ show opts) *>
    runSystem (mkLoadModuleCmd mode name opts) ("insmod " ++ name ++ " error.")



setupDevice :: Maybe Int -> Device -> IO ()
setupDevice queues (Device dev speed channels fctrl opts) = do

    putStrBoldLn $ "Activating " ++ dev ++ "..."

    runSystem (ifconfig ++ " " ++ dev ++ " up") "ifconfig error!"

    threadDelay defaultDelay

    if isJust speed
       then do
            let s = fromJust speed
            putStrBoldLn $ "Setting speed (" ++ show s ++ ") for " ++ dev ++ "..."
            runSystem (ethtool ++ " -s " ++ dev ++ " speed " ++ show s ++ " duplex full autoneg off") "ethtool: set speed error!"
       else
            putStrBoldLn ("Auto-negotiating speed for " ++ dev ++ "...") *>
            runSystem (ethtool ++ " -s " ++ dev ++ " duplex full autoneg off") "ethtool: auto-negotiation error!"

    threadDelay defaultDelay

    when (isJust fctrl) $ do
        let t = fromJust fctrl
        if t then putStrBoldLn ("Disabling flow control for " ++ dev ++ "...") *>
                    runSystem (ethtool ++ " -A " ++ dev ++ " rx on tx on") "ethtool: enabling flow-ctrl error!"
             else putStrBoldLn ("Enabling flow control for " ++ dev ++ "...") *>
                    runSystem (ethtool ++ " -A " ++ dev ++ " autoneg off rx off tx off") "ethtool: disabling flow-ctrl error!"

    threadDelay defaultDelay

    when (isJust queues || isJust channels) $ do
        let c = fromJust (channels <|> queues)
        putStrBoldLn $ "Setting channels to " ++ show c ++ "..."
        runSystem (ethtool ++ " -L " ++ dev ++ " combined " ++ show c) ""

    forM_ opts $ \(opt, arg) ->
        threadDelay defaultDelay *>
        runSystem (ethtool ++ " " ++ opt ++ " " ++ dev ++ " " ++ arg) ("ethtool:" ++ opt ++ " error!")


setupPCI :: [(String, String)] -> IO ()
setupPCI = mapM_ (\(a, v) -> runSystem (setpci ++ " -v -d " ++ show a ++ " " ++ show v) (setpci ++ ": error!"))


getDevices ::  Config -> [String]
getDevices conf =
    map devname (concatMap devices (drivers conf))


setupIRQAffinity :: Int -> [Int] -> [String] -> [String] -> IO ()
setupIRQAffinity fc excl algs devs = do
    let excl_opt = unwords (map (\n -> " -e " ++ show n) excl)
    let affinity = zip algs (tails devs)
    unless (null affinity) $
        forM_ affinity $ \(alg, devs') ->
            runSystem ("pfq-affinity -f " ++ show fc  ++ " " ++ excl_opt ++ " -a " ++ alg ++ " -m TxRx " ++ unwords devs') "pfq-affinity error!"


runSystem :: String -> String -> IO ()
runSystem cmd errmsg =
    putStrLn ("-> " ++ cmd) *>
        system cmd >>= \ec -> unless (null errmsg || ec == ExitSuccess) $ throw (SystemError (red ++ errmsg ++ reset) ec)


putStrBoldLn :: String -> IO ()
putStrBoldLn msg = putStrLn $ bold ++ msg ++ reset


possiblyReadFile :: FilePath -> IO String
possiblyReadFile file = do
    res <- tryIOError $ readFile file
    case res of
        Left _    -> return ""
        Right str -> length str `seq` return str

