{-# LANGUAGE Arrows, TransformListComp, RecordWildCards #-}

module Main where

import Nettle.FRPControl.NettleSF
import Nettle.FRPControl.NetInfo
import Nettle.IPv4.IPPacket
import Nettle.IPv4.IPAddress
import Nettle.OpenFlow.Match
import Nettle.OpenFlow.Packet
import Nettle.OpenFlow.Action
import Nettle.OpenFlow.Statistics
import Nettle.OpenFlow.Switch
import Nettle.Ethernet.EthernetFrame
import Data.Maybe (fromMaybe)
import qualified Data.Map as Map

{- 
A learning switch, which installs flow table entries at the switches to 
to avoid switch-to-controller messages.
-}

-- Warning: Only works for loop-free networks; On loopy networks, a single packet will loop indefinitely 
-- through the network.

main :: IO ()
main = simpleNettleDriver 2525 (controller)

monitorPolicy = NetworkMonitorPolicy { switchFeaturesRefreshPeriod = 20, -- seconds
                                       portStatisticsRefreshPeriod = 20   -- seconds
                                     } 

controller :: SF (Event (SwitchID, SwitchMessage)) (Event SwitchCommand, Event String)
controller = proc i -> do 
  netInfoCmdE <- networkInfoRequester monitorPolicy  -< i
  initCmd <- switchInitializer' -< i
  hostDirMap <- hostDirectionsSF -< i
  let tableMods = mapFilterE (updateTableCommands hostDirMap) (packetInE i)
  returnA -< (mergeEventsBy (<+>) [netInfoCmdE, initCmd, tableMods], noEvent)

updateTableCommands hostDirMap (sw, pktInfo@PacketInfo {..}) = 
  case packetInFrame pktInfo of 
    Left msg       -> Nothing
    Right ethFrame -> 
      let pred = fromMatch (frameToExactMatch receivedOnPort ethFrame)
          s    = sourceAddress ethFrame
          r    = destAddress ethFrame
          cmd = do ps <- Map.lookup (sw,s) hostDirMap
                   pr <- Map.lookup (sw,r) hostDirMap             
                   case bufferID of
                     Nothing -> return (addFlowRule (((pred ==> sendOnPort pr) `expiringAfter` 240) `withPriority` 1) sw
                                        <+> sendPacket sw (receivedPacketOut pktInfo flood)
                                       )
                     Just bufid -> return (addFlowRule' (((pred ==> sendOnPort pr) `expiringAfter` 240) `withPriority` 1) bufid sw)
      in case cmd of 
        Nothing -> Just (sendPacket sw (receivedPacketOut pktInfo flood))
        Just cmd' -> Just cmd'
                
{-
updateTableCommands hostDirMap (sw,pktInfo) =
  case packetInFrame pktInfo of 
    Left msg       -> Nothing
    Right ethFrame -> 
       let s = sourceAddress ethFrame
           r = destAddress ethFrame
       in do ps <- Map.lookup (sw,s) hostDirMap
             pr <- Map.lookup (sw,r) hostDirMap             
             return (deleteFlowRules (ethSourceDestAre s r <|> ethSourceDestAre r s) sw <+>
                     addFlowRules [(1, flowFromTo s ps r pr `expiringAfter` 240) , 
                                   (1, flowFromTo r pr s ps `expiringAfter` 240)] sw)
    where flowFromTo s ps r pr = inPortIs ps <&> ethSourceDestAre s r ==> sendOnPort pr 
-}
switchInitializer' :: SF (Event (SwitchID, SwitchMessage)) (Event SwitchCommand)
switchInitializer' = proc i -> do
  returnA -< liftE (\(swid,_) -> clearTables swid <+> 
                                 addFlowRule ((arp ==> flood) `withPriority` 1) swid) (arrivalE i)  

packetSender :: SF (Event (SwitchID, SwitchMessage)) (Event SwitchCommand)
packetSender = proc i -> do 
  returnA -< liftE (\(switchID, pktIn) -> sendPacket switchID (receivedPacketOut pktIn flood)) (packetInE i)