{-# LANGUAGE RecordWildCards #-}

module Main where

import Nettle.Servers.MultiplexedTCPServer
import Nettle.Servers.TCPServer
import SimpleImperativeIONetworkControl
import Nettle.OpenFlow.Messages as M
import Nettle.OpenFlow.Switch
import Nettle.OpenFlow.Port
import Nettle.OpenFlow.Packet
import Nettle.OpenFlow.Action
import Nettle.OpenFlow.FlowTable
import Nettle.OpenFlow.Match
import Nettle.Ethernet.EthernetAddress
import Nettle.Ethernet.EthernetFrame
import Nettle.IPv4.IPPacket
import Control.Monad.State
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Binary
import Data.Binary.Get
import Control.Monad.Error


type ControlState = (HostTable, SwitchState)
type SwitchState = Map.Map SockAddr SwitchFeatures 

{- Functions for getting and modifying the switch state. -}
switchFeatures :: SockAddr -> NetController ControlState (Maybe SwitchFeatures)
switchFeatures addr = gets snd >>= return . Map.lookup addr

addSwitchFeatures :: SockAddr -> SwitchFeatures -> ControlState -> ControlState
addSwitchFeatures addr sfr (htbl,stbl) = (htbl, Map.insert addr sfr stbl)

deleteSwitch :: SockAddr -> ControlState -> ControlState
deleteSwitch addr (htbl,stbl) = (htbl, Map.delete addr stbl)


{- 
The host table mapps hosts to their known locations, i.e. to physical ports at OpenFlow switches. 
-}

type HostTable = Map.Map (EthernetAddress,SockAddr) PortID

newHostTable = Map.empty

learn ::  EthernetAddress -> SockAddr -> PortID -> HostTable -> HostTable
learn hostAddr switchAddr portID = Map.insert (hostAddr, switchAddr) portID

hostLookup ::  EthernetAddress -> SockAddr -> HostTable -> Maybe PortID
hostLookup hostAddr switchAddr htbl = Map.lookup (hostAddr, switchAddr) htbl

{- Convenience functions for manipulating state -}

hostTableM :: NetController ControlState HostTable
hostTableM = gets fst

learnM :: EthernetAddress -> SockAddr -> PortID -> NetController ControlState ()
learnM addr dpid inPort = hostTableM >>= updateHostTableM . learn addr dpid inPort 

hostLookupM :: EthernetAddress -> SockAddr -> NetController ControlState (Maybe PortID)
hostLookupM hostAddr switchAddr = hostTableM >>= return . hostLookup hostAddr switchAddr

updateHostTableM :: HostTable -> NetController ControlState ()
updateHostTableM htbl = modify (\(_,swState) -> (htbl,swState))



{- The controller -}

main = runNetController 2525 controller (newHostTable, Map.empty)

controller :: NetController ControlState ()
controller = mainHandler Set.empty
    where mainHandler needHellos = 
              do e <- waitForEvent 
                 case e of
                   ConnectionEstablished addr -> 
                       do sendMessage addr 0 CSHello
                          mainHandler (Set.insert addr needHellos)
                   ConnectionTerminated addr e -> 
                       do modify (deleteSwitch addr)
                          mainHandler (Set.delete addr needHellos)
                   PeerMessage addr (xid, msg) -> 
                       if addr `Set.member` needHellos
                       then case msg of 
                              SCHello -> do sendMessage addr 0 FeaturesRequest
                                            mainHandler (Set.delete addr needHellos) 
                              _       -> do liftIO $ putStrLn ("Expected to receive Hello from " ++ show addr ++ ", but received " ++ show msg)
                       else do case msg of 
                                 SCHello              -> liftIO $ putStrLn ("Received unexpected hello from " ++ show addr ++ ".")
                                 SCEchoRequest bytes  -> sendMessage addr 0 (CSEchoReply bytes)
                                 Features sfr         -> modify (addSwitchFeatures addr sfr)
                                 PacketIn pktInRecord -> handlePacketIn addr pktInRecord
                                 _                    -> liftIO $ putStrLn ("Unhandled message: " ++ show msg)
                               mainHandler needHellos

handlePacketIn addr pktIn@(PacketInfo {..}) = 
   case runGetE getEthernetFrame packetData of
    Left errorMessage -> liftIO $ putStrLn ("Failed to parse ethernet frame: " ++ errorMessage)
    Right ethernetFrame ->
      do learnM (sourceAddress ethernetFrame) addr receivedOnPort
         if isReserved (sourceAddress ethernetFrame)
           then sendMessage addr 0 $ M.FlowMod $ AddFlow { match             = frameToExactMatch receivedOnPort ethernetFrame
                                                         , actions           = []
                                                         , priority          = 0
                                                         , idleTimeOut       = ExpireAfter 60
                                                         , hardTimeOut       = Permanent
                                                         , applyToPacket     = bufferID
                                                         , overlapAllowed    = True
                                                         , notifyWhenRemoved = False
                                                         , cookie            = 0
                                                         }
           else do mport <- hostLookupM (destAddress ethernetFrame) addr
                   case mport of 
                     Nothing   -> sendMessage addr 0 (M.PacketOut $ receivedPacketOut pktIn flood)
                     Just lprt -> do liftIO $ putStrLn $ "at addr: " ++ show addr ++ " forward to port " ++ show lprt
                                     sendMessage addr 0 $ M.FlowMod $ AddFlow { match             = frameToExactMatch receivedOnPort ethernetFrame
                                                                              , actions           = sendOnPort lprt
                                                                              , priority          = 0
                                                                              , idleTimeOut       = ExpireAfter 60
                                                                              , hardTimeOut       = Permanent
                                                                              , applyToPacket     = bufferID
                                                                              , overlapAllowed    = True
                                                                              , notifyWhenRemoved = False
                                                                              , cookie            = 0
                                                                              }
                                     sendMessage addr 0 $ M.PacketOut $ receivedPacketOut pktIn (sendOnPort lprt)