{-|
Module      : MPTCP
Description : Implementation of mptcp netlink path manager
Maintainer  : matt
Stability   : testing
Portability : Linux

OverloadedStrings allows Aeson to convert
-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
module Net.Mptcp
where

import Net.SockDiag ()
import Control.Exception (assert)

import Data.Word (Word8, Word16, Word32)
import qualified Data.Map as Map
import System.Linux.Netlink hiding (makeSocket)
-- import System.Linux.Netlink (query, Packet(..))
import System.Linux.Netlink.GeNetlink
import System.Linux.Netlink.Constants
-- import System.Linux.Netlink.GeNetlink.Control
import Data.ByteString (ByteString)
import Data.Maybe (fromJust)

import Data.Serialize.Get
import Data.Serialize.Put

import Data.Bits ((.|.))

import Data.List ()
import Debug.Trace
import Control.Concurrent ()
import Net.IP
import Net.Tcp
import Net.IPAddress
import Net.IPv4
import Net.IPv6
import Net.Mptcp.Constants
-- import Data.Text
-- import Data.Set ()
import qualified Data.Set as Set
-- in package Unique-0.4.7.6
-- import Data.List.Unique
import Data.Aeson
import GHC.Generics
-- import Data.Default

-- how can I retreive the word16 without pattern matching ?
data MptcpSocket = MptcpSocket NetlinkSocket Word16
instance Show MptcpSocket where
  show sock = let (MptcpSocket nlSock fid) = sock in ("Mptcp netlink socket: " ++ show fid)

type MptcpPacket = GenlPacket NoData

-- |Same as SockDiagMetrics
-- data SubflowWithMetrics = SubflowWithMetrics {
--   subflowSubflow :: TcpConnection
--     -- for now let's retain DiagTcpInfo  only
--   , metrics :: [SockDiagExtension]
-- }

-- |Data to hold MPTCP level information
-- TODO use Data.Set
data MptcpConnection = MptcpConnection {
  connectionToken :: MptcpToken
  -- use SubflowWithMetrics instead ?!
  -- , subflows :: Set.Set [TcpConnection]
  , subflows :: Set.Set TcpConnection
  -- TODO convert to Data.Set ?
  -- , localIds :: [Word8]  -- ^ Announced addresses
  -- , remoteIds :: [Word8]  -- ^ Announced addresses
  , localIds :: Set.Set Word8  -- ^ Announced addresses
  , remoteIds :: Set.Set Word8   -- ^ Announced addresses

  -- Might be reworked/moved in an Enriched/Tracker structure afterwards
  -- , tcpMetrics :: Maybe [SockDiagExtension]  -- ^Metrics retrieved from kernel
  , get_caps_prog :: FilePath
} deriving (Show, Generic)

-- | Remote port
data RemoteId = RemoteId {

  remoteAddress :: IP
  , remotePort :: Word16
}


-- data MptcpAttributes = Map.Map MptcpAttr MptcpAttribute


remoteIdFromAttributes :: Attributes -> RemoteId
remoteIdFromAttributes attrs = let
    (SubflowDestPort dport) = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_DPORT attrs
    (SubflowFamily family) = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_FAMILY attrs
    SubflowDestAddress destIp = ipFromAttributes False attrs

    -- (SubflowDestPort dport) = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_DPORT attrs
  in
    RemoteId destIp dport

-- we don't really care
instance FromJSON MptcpConnection

-- | export to the format expected by mptcpnumerics
-- could be automatically generated ?
-- toJSON :: MptcpConnection -> Value
instance ToJSON MptcpConnection where
  toJSON mptcpConn = object
    [ "name" .= toJSON (show $ connectionToken mptcpConn)
    , "sender" .= object [
          -- TODO here we could read from sysctl ? or use another SockDiagExtension
          "snd_buffer" .= toJSON (40 :: Int)
          , "capabilities" .= object []
        ]
    , "capabilities" .= object ([])
    -- TODO generated somewhere else
    -- , "subflows" .= object ([])
    ]


--
-- instance ToJSON SubflowWithMetrics where
--   toJSON sf = object [
--     (pack (nameFromTcpConnection $ subflowSubflow sf) .= object [

--       "cwnd" .= toJSON (20 :: Int),
--       -- for now hardcode mss ? we could set it to one to make
--       "mss" .= toJSON (1500 :: Int),
--       "var" .= toJSON (10 :: Int),
--       "fowd" .= toJSON (10 :: Int),
--       "bowd" .= toJSON (10 :: Int),
--       "loss" .= toJSON (0.5 :: Float)
--       -- This is an user preference, that should be pushed when calling mptcpnumerics
--       -- , "contribution": 0.5
--     ])
--     ]


-- |Adds a subflow to the connection
-- use sets instead ?
-- TODO compose with mptcpConnAddLocalId
mptcpConnAddSubflow :: MptcpConnection -> TcpConnection -> MptcpConnection
mptcpConnAddSubflow mptcpConn sf =
  -- trace ("Adding subflow" ++ show sf)
    mptcpConnAddLocalId
        (mptcpConnAddRemoteId
            -- (mptcpConn { subflows = sf : subflows mptcpConn })
            (mptcpConn { subflows = Set.insert sf (subflows mptcpConn) })
            (remoteId sf)
        )
        (localId sf)

    -- , localIds = Set.insert (localId sf) (localIds mptcpConn)
    -- , remoteIds = Set.insert (remoteId sf) (remoteIds mptcpConn)
  -- }


-- |Add local id
mptcpConnAddLocalId :: MptcpConnection
                       -> Word8 -- ^Local id to add
                       -> MptcpConnection
mptcpConnAddLocalId con locId = con { localIds = Set.insert (locId) (localIds con) }


-- |Add remote id
-- Le remote id doit etre associee a une addresse / TODO test protocol
mptcpConnAddRemoteId :: MptcpConnection
                       -> Word8 -- ^Remote id to add
                       -> MptcpConnection
mptcpConnAddRemoteId con remId = con { localIds = Set.insert (remId) (remoteIds con) }

-- |Remove subflow from an MPTCP connection
mptcpConnRemoveSubflow :: MptcpConnection -> TcpConnection -> MptcpConnection
mptcpConnRemoveSubflow con sf = con {
  subflows = Set.delete sf (subflows con)
  -- TODO remove associated local/remote Id ?
}


getPort :: ByteString -> Word16
getPort val =
  case (runGet getWord16host val) of
    Left _ -> 0
    Right port -> port


--
-- The message type/ flag / sequence number / pid  (0 => from the kernel)
-- https://elixir.bootlin.com/linux/latest/source/include/uapi/linux/netlink.h#L54
fixHeader :: MptcpSocket -> Bool -> MptcpPacket -> MptcpPacket
fixHeader (MptcpSocket _ fid) dump pkt = let
    myHeader = Header 0 (flags .|. fNLM_F_ACK) 0 0
    flags = if dump then fNLM_F_REQUEST .|. fNLM_F_MATCH .|. fNLM_F_ROOT else fNLM_F_REQUEST
  in
    pkt { packetHeader = myHeader }


{-|
  Generates an Mptcp netlink request
TODO we could fake the Word16/Flag and 
-}
genMptcpRequest :: Word16 -- ^the family id
                -> MptcpGenlEvent -- ^The MPTCP command
                -> Bool           -- ^Dump answer (returns EOPNOTSUPP if not possible)
                -- -> Attributes
                -> [MptcpAttribute]
                -> MptcpPacket
genMptcpRequest fid cmd dump attrs =
  let
    myHeader = Header (fromIntegral fid) (flags .|. fNLM_F_ACK) 0 0
    geheader = GenlHeader word8Cmd mptcpGenlVer
    flags = if dump then fNLM_F_REQUEST .|. fNLM_F_MATCH .|. fNLM_F_ROOT else fNLM_F_REQUEST
    word8Cmd = fromIntegral (fromEnum cmd) :: Word8

    pkt = Packet myHeader (GenlData geheader NoData) (mptcpListToAttributes attrs)
    -- TODO run an assert on the list filter
    hasTokenAttr = Prelude.any (isAttribute (MptcpAttrToken 0)) attrs
  in
    assert hasTokenAttr pkt

readToken :: Maybe ByteString -> MptcpToken
readToken maybeVal = case maybeVal of
  Nothing -> error "Missing token"
  Just val -> case ( runGet getWord32host val) of
    Left err -> error "could not decode"
    Right mptcpToken -> mptcpToken

-- LocId => Word8
readLocId :: Maybe ByteString -> LocId
readLocId maybeVal = case maybeVal of
  Nothing -> error "Missing locator id"
  Just val -> case runGet getWord8 val of
    -- TODO generate an error here !
    Left _ -> error "Could not get locId !!"
    Right locId -> locId
  -- runGet getWord8 val

-- doDumpLoop :: MyState -> IO MyState
-- doDumpLoop myState = do
--     let (MptcpSocket simpleSock fid) = socket myState
--     results <- recvOne' simpleSock ::  IO [Either String MptcpPacket]
--     -- TODO retrieve packets
--     mapM_ (inspectResult myState) results
--     newState <- doDumpLoop myState
--     return newState


-- data MptcpAttributes = MptcpAttributes {
--     connToken :: Word32
--     , localLocatorID :: Maybe Word8
--     , remoteLocatorID :: Maybe Word8
--     , family :: Word16 -- Remove ?
--     -- |Pointer to the Attributes map used to build this struct. This is purely
--     -- |for forward compat, please file a feature report if you have to use this.
--     , staSelf       :: Attributes
-- } deriving (Show, Eq, Read)

-- Wouldn't it be easier to work with ?
-- data MptcpEvent = NewConnection {
-- }


-- |Represents every possible setting sent/received on the netlink channel
data MptcpAttribute =
    MptcpAttrToken MptcpToken |
    -- v4 or v6, AddressFamily is a netlink def
    SubflowFamily AddressFamily | -- ^ should be Word16 too
    -- remote/local ?
    RemoteLocatorId Word8 |
    LocalLocatorId Word8 |
    SubflowSourceAddress IP |
    SubflowDestAddress IP |
    SubflowSourcePort Word16 |
    SubflowDestPort Word16 |
    SubflowMaxCwnd Word32 |
    SubflowBackup Word8 |
    SubflowInterface Word32
    deriving (Show, Eq)

type MptcpToken = Word32
type LocId    = Word8

-- inspired by netlink cATA :: CtrlAttribute -> (Int, ByteString)
attrToPair :: MptcpAttribute -> (Int, ByteString)
attrToPair (MptcpAttrToken token) = (fromEnum MPTCP_ATTR_TOKEN, runPut $ putWord32host token)
attrToPair (RemoteLocatorId loc) = (fromEnum MPTCP_ATTR_REM_ID, runPut $ putWord8 loc)
attrToPair (LocalLocatorId loc) = (fromEnum MPTCP_ATTR_LOC_ID, runPut $ putWord8 loc)
attrToPair (SubflowFamily fam) = let
        fam8 = (fromIntegral $ fromEnum fam) :: Word16
    in (fromEnum MPTCP_ATTR_FAMILY, runPut $ putWord16host fam8)

attrToPair ( SubflowInterface idx) = (fromEnum MPTCP_ATTR_IF_IDX, runPut $ putWord32host idx)
attrToPair ( SubflowSourcePort port) = (fromEnum MPTCP_ATTR_SPORT, runPut $ putWord16host port)
attrToPair ( SubflowDestPort port) = (fromEnum MPTCP_ATTR_DPORT, runPut $ putWord16host port)
attrToPair ( SubflowMaxCwnd limit) = (fromEnum MPTCP_ATTR_CWND, runPut $ putWord32host limit)
attrToPair ( SubflowBackup prio) = (fromEnum MPTCP_ATTR_BACKUP, runPut $ putWord8 prio)
-- TODO should depend on the ip putWord32be w32
attrToPair ( SubflowSourceAddress addr) =
  case_ (genV4SubflowAddress MPTCP_ATTR_SADDR4) (genV6SubflowAddress MPTCP_ATTR_SADDR6) addr
attrToPair ( SubflowDestAddress addr) =
  case_ (genV4SubflowAddress MPTCP_ATTR_DADDR4) (genV6SubflowAddress MPTCP_ATTR_DADDR6) addr


genV4SubflowAddress :: MptcpAttr -> IPv4 -> (Int, ByteString)
genV4SubflowAddress attr ip = (fromEnum attr, runPut $ putWord32be w32)
  where
    w32 = getIPv4 ip

genV6SubflowAddress :: MptcpAttr -> IPv6 -> (Int, ByteString)
genV6SubflowAddress addr = undefined

mptcpListToAttributes :: [MptcpAttribute] -> Attributes
mptcpListToAttributes attrs = Map.fromList $Prelude.map attrToPair attrs


-- |Retreive IP
-- TODO could check/use addressfamily as well
ipFromAttributes :: Bool  -- ^True if source
                    -> Attributes -> MptcpAttribute
ipFromAttributes True attrs =
    case makeAttributeFromMaybe MPTCP_ATTR_SADDR4 attrs of
      Just ip -> ip
      Nothing -> case makeAttributeFromMaybe MPTCP_ATTR_SADDR6 attrs of
        Just ip -> ip
        Nothing -> error "could not get the src IP"

ipFromAttributes False attrs =
    case makeAttributeFromMaybe MPTCP_ATTR_DADDR4 attrs of
      Just ip -> ip
      Nothing -> case makeAttributeFromMaybe MPTCP_ATTR_DADDR6 attrs of
        Just ip -> ip
        Nothing -> error "could not get dest IP"

-- mptcpAttributesToMap :: [MptcpAttribute] -> Attributes
-- mptcpAttributesToMap attrs =
--   Map.fromList $map mptcpAttributeToTuple attrs

-- |Converts / should be a maybe ?
-- TODO simplify
subflowFromAttributes :: Attributes -> TcpConnection
subflowFromAttributes attrs =
  let
    -- expects a ByteString
    SubflowSourcePort sport = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_SPORT attrs
    SubflowDestPort dport = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_DPORT attrs
    SubflowSourceAddress _srcIp =  ipFromAttributes True attrs
    SubflowDestAddress _dstIp = ipFromAttributes False attrs
    LocalLocatorId lid = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_LOC_ID attrs
    RemoteLocatorId rid = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_REM_ID attrs
    SubflowInterface intfId = fromJust $ makeAttributeFromMaybe MPTCP_ATTR_IF_IDX attrs
    -- sfFamily = getPort $ fromJust (Map.lookup (fromEnum MPTCP_ATTR_FAMILY) attrs)
    prio = Nothing   -- (SubflowPriority N)
  in
    -- TODO fix sfFamily
    TcpConnection _srcIp _dstIp sport dport prio lid rid (Just intfId)


-- TODO prefix with 'e' for enum
-- Map.lookup (fromEnum attr) m
-- getAttribute :: MptcpAttr -> Attributes -> Maybe MptcpAttribute
-- getAttribute attr m
--     | attr == MPTCP_ATTR_TOKEN = Nothing
--     | otherwise = Nothing

-- getAttribute :: (Int, ByteString) -> CtrlAttribute
-- getAttribute (i, x) = fromMaybe (CTRL_ATTR_UNKNOWN i x) $makeAttribute i x

-- getW16 :: ByteString -> Maybe Word16
-- getW16 x = e2M (runGet g16 x)

-- getW32 :: ByteString -> Maybe Word32
-- getW32 x = e2M (runGet g32 x)

-- "either2Maybe"
e2M :: Either a b -> Maybe b
e2M (Right x) = Just x
e2M _ = Nothing

-- un peu bidon
-- convertAttributesIntoList :: Attributes -> [MptcpAttribute]
-- convertAttributesIntoList attrs = let
--       customMakeAttribute key val accum = [fromJust (makeAttribute key val)] ++ accum
--   in 
--       -- foldrWithKey." => (k -> a -> b -> b) -> b -> Map k a -> b
--       Map.foldrWithKey (customMakeAttribute) [] attrs

convertAttributesIntoMap :: Attributes -> Map.Map MptcpAttr MptcpAttribute
convertAttributesIntoMap attrs = let
-- traverseWithKey
-- fromJust
      customFn k val = fromJust (makeAttribute k val)
      -- fromJust $ makeAttribute
      -- (k -> a -> b) -> Map k a -> Map k b
      newMap = Map.mapWithKey (customFn) attrs
  in
      Map.mapKeys (toEnum) newMap

-- TODO rename fromMap
makeAttributeFromMaybe :: MptcpAttr -> Attributes -> Maybe MptcpAttribute
makeAttributeFromMaybe attrType attrs =
  let res = Map.lookup (fromEnum attrType) attrs in
  case res of
    Nothing -> error $ "Could not build attr " ++ show attrType
    Just bytestring -> makeAttribute (fromEnum attrType) bytestring

-- | Builds an MptcpAttribute from
makeAttribute :: Int -- ^ MPTCP_ATTR_TOKEN value
                  -> ByteString
                  -> Maybe MptcpAttribute
makeAttribute i val =
  case toEnum i of
    MPTCP_ATTR_TOKEN -> Just (MptcpAttrToken $ readToken $ Just val)
    -- TODO fix
    MPTCP_ATTR_FAMILY ->
        case runGet getWord16host val of
          -- assert it's eAF_INET or eAF_INET6
          Right x -> Just $ SubflowFamily (toEnum ( fromIntegral x :: Int))
          _ -> Nothing
    MPTCP_ATTR_SADDR4 -> SubflowSourceAddress <$> fromIPv4 <$> e2M ( getIPv4FromByteString val)
    MPTCP_ATTR_DADDR4 -> SubflowDestAddress <$> fromIPv4 <$> e2M (getIPv4FromByteString val)
    MPTCP_ATTR_SADDR6 -> SubflowSourceAddress <$> fromIPv6 <$> e2M (getIPv6FromByteString val)
    MPTCP_ATTR_DADDR6 -> SubflowDestAddress <$> fromIPv6 <$> e2M (getIPv6FromByteString val)
    MPTCP_ATTR_SPORT -> SubflowSourcePort <$> port where port = e2M $ runGet getWord16host val
    MPTCP_ATTR_DPORT -> SubflowDestPort <$> port where port = e2M $ runGet getWord16host val
    MPTCP_ATTR_LOC_ID -> Just (LocalLocatorId $ readLocId $ Just val )
    MPTCP_ATTR_REM_ID -> Just (RemoteLocatorId $ readLocId $ Just val )
    MPTCP_ATTR_IF_IDX -> trace ("if_idx: " ++ show val) (
             case runGet getWord32be val of
                Right x -> Just $ SubflowInterface x
                _ -> Nothing)
    -- backup is u8
    MPTCP_ATTR_BACKUP -> Just (SubflowBackup $ readLocId $ Just val )
    MPTCP_ATTR_ERROR -> trace "makeAttribute ERROR" Nothing
    MPTCP_ATTR_TIMEOUT -> undefined
    MPTCP_ATTR_CWND -> undefined
    MPTCP_ATTR_FLAGS -> trace "makeAttribute ATTR_FLAGS" Nothing
    MPTCP_ATTR_UNSPEC -> undefined


dumpAttribute :: Int -> ByteString -> String
dumpAttribute attrId value =
  show $ makeAttribute attrId value

checkIfSocketExistsPkt :: Word16 -> [MptcpAttribute]  -> MptcpPacket
checkIfSocketExistsPkt fid attributes =
    genMptcpRequest fid MPTCP_CMD_EXIST True attributes

-- https://stackoverflow.com/questions/47861648/a-general-way-of-comparing-constructors-of-two-terms-in-haskell?noredirect=1&lq=1
-- attrToPair ( SubflowSourcePort port) = (fromEnum MPTCP_ATTR_SPORT, runPut $ putWord8 loc)
isAttribute :: MptcpAttribute -- ^ to compare with
               -> MptcpAttribute -- ^to compare to
               -> Bool
isAttribute ref toCompare = fst (attrToPair toCompare) == fst (attrToPair ref)

-- create a fake LocalLocatorId
hasLocAddr :: [MptcpAttribute] -> Bool
hasLocAddr attrs = Prelude.any (isAttribute (LocalLocatorId 0)) attrs

hasFamily :: [MptcpAttribute] -> Bool
hasFamily = Prelude.any (isAttribute (SubflowFamily eAF_INET))

-- need to prepare a request
-- type GenlPacket a = Packet (GenlData a)
-- REQUIRES: LOC_ID / TOKEN
-- TODO pass TcpConnection 
resetConnectionPkt :: MptcpSocket -> [MptcpAttribute] -> MptcpPacket
resetConnectionPkt (MptcpSocket sock fid) attrs = let
    cmd = MPTCP_CMD_REMOVE
  in
    assert (hasLocAddr attrs) $ genMptcpRequest fid MPTCP_CMD_REMOVE False attrs


-- TODO map an IP to an interface
-- pass token ?
-- TODO can we retreive ips as well ?!
subflowAttrs :: TcpConnection -> [MptcpAttribute]
subflowAttrs con = [
    LocalLocatorId $ localId con
    , RemoteLocatorId $ remoteId con
    -- TODO adapt
    , SubflowFamily $ getAddressFamily (dstIp con)
    , SubflowDestAddress $ dstIp con
    , SubflowDestPort $ dstPort con
    -- should fail if doesn't exist
    , SubflowInterface $ fromJust $ subflowInterface con
    -- https://github.com/multipath-tcp/mptcp/issues/338
    , SubflowSourceAddress $ srcIp con
    , SubflowSourcePort $ srcPort con
  ]

-- |Generate a request to create a new subflow
-- TODO
capCwndPkt :: MptcpSocket -> MptcpConnection
              -> Word32  -- ^Limit to apply to congestion window
              -> TcpConnection -> MptcpPacket
capCwndPkt (MptcpSocket sock fid) mptcpCon limit sf =
    assert (hasFamily attrs) pkt
    where
        oldPkt = genMptcpRequest fid MPTCP_CMD_SND_CLAMP_WINDOW False attrs
        -- TODO maybe cleanup interface after that: try to update Header
        -- messageSeqNum =
        -- messagePID
        pkt = oldPkt { packetHeader = (packetHeader oldPkt) { messagePID = 42 } }
        attrs = connectionAttrs mptcpCon
              ++ [ SubflowMaxCwnd limit ]
              ++ subflowAttrs sf


connectionAttrs :: MptcpConnection -> [MptcpAttribute]
connectionAttrs con = [ MptcpAttrToken $ connectionToken con ]

-- sport/backup/intf are optional
-- family /loc id/remid/daddr/dport
-- TODO pass new subflow ?
-- should get rid of MptcpSocket  [MptcpAttribute] ->
newSubflowPkt :: MptcpSocket -> MptcpConnection -> TcpConnection -> MptcpPacket
newSubflowPkt (MptcpSocket sock fid) mptcpCon sf = let
    cmd = MPTCP_CMD_SUB_CREATE
    attrs = connectionAttrs mptcpCon ++ subflowAttrs sf
    pkt = genMptcpRequest fid MPTCP_CMD_SUB_CREATE False attrs
  in
    assert (hasFamily attrs) pkt

-- announceLocalIdPkt ::
-- announceLocalIdPkt =