package Bus(mkBus, Bus(..)) where

import Types
import BusTypes
import TagEngine
import Vector
import Util
import Arbiter
import FIFO
import FIFOF
import SpecialFIFOs
import Printf

busRequestToAddr :: BusRequest -> Addr
busRequestToAddr req = case req of
    BusReadRequest (ReadRequest addr _) -> addr
    BusWriteRequest (WriteRequest addr _) -> addr

-- Create a Bus Module that supports multiple clients and servers
-- submitting requests and simultaneously returning responses.
-- Responses can be consumed by clients out of order as all client
-- submitted requests are tagged - and servers keep that tag
-- when responding.
mkBus :: (Add n (TLog numServers) (TLog (TAdd numServers 1)))
      => (Addr -> Maybe (MkServerIdx numServers))
      -> Module (Bus inFlightTransactions numClients numServers)
mkBus serverMap = do
    -- Tag engines for each client to manage transaction tags
    tagEngineByClientVec :: Vector numClients (TagEngine inFlightTransactions)
    tagEngineByClientVec <- replicateM mkTagEngine

    -- There are `numClients` clients, each of which needs its own arbiter as
    -- there are up to `numServer` servers that may wish to submit a response
    -- to a given client. Furthermore the rule that routes client requests to
    -- servers makes for another potential requestor as it may determine that
    -- a request is unmappable and instead opt to form and submit a
    -- `BusError UnMapped` response directly to a client response arbiter. Thus
    -- we must arbit between a total of `numServers + 1` requestors.
    responseArbiterByClient :: Vector numClients (Arbiter.Arbiter_IFC (TAdd numServers 1))
    responseArbiterByClient <- replicateM (mkArbiter False)

    -- There are `numServer` servers, each of which needs its own arbiter as
    -- there are up to `numClient` clients that may wish to submit a response
    -- to a given server.
    requestArbiterByServer :: Vector numServers (Arbiter.Arbiter_IFC numClients)
    requestArbiterByServer <- replicateM (mkArbiter False)

    dummyVar :: Reg(Bool)
    dummyVar <- mkReg False

    -- Queues to hold requests from clients
    clientRequestQueues :: Vector numClients (FIFOF (TaggedBusRequest inFlightTransactions))
    clientRequestQueues <- replicateM (mkSizedBypassFIFOF (valueOf inFlightTransactions))

    -- Queues to hold responses to clients
    clientResponseQueues :: Vector numClients (FIFOF (TaggedBusResponse inFlightTransactions))
    clientResponseQueues <- replicateM (mkSizedBypassFIFOF (valueOf inFlightTransactions))

    -- The following two vectors of FIFOs make it easier to push/pull data to/from internal
    -- server methods:
    consumeRequestQueues :: Vector numServers (
                                        FIFOF (
                                            MkClientTagType numClients,
                                            TaggedBusRequest inFlightTransactions
                                            )
                                        )
    consumeRequestQueues <- replicateM mkBypassFIFOF

    submitResponseQueues :: Vector numServers (
                                        FIFOF (
                                            MkClientTagType numClients,
                                            TaggedBusResponse inFlightTransactions
                                            )
                                        )
    submitResponseQueues <- replicateM mkBypassFIFOF

    let clientRules :: Vector numClients (Rules)
        clientRules = genWith $ \clientIdx ->
            let
                selectedClientRequestQueue :: FIFOF (TaggedBusRequest inFlightTransactions)
                selectedClientRequestQueue = (select clientRequestQueues clientIdx)
            in
                rules
                    (sprintf "client[%d] route request" clientIdx): when True ==> do
                        let
                            clientRequest :: TaggedBusRequest inFlightTransactions
                            clientRequest = selectedClientRequestQueue.first

                            targetAddr :: Addr = busRequestToAddr |> clientRequest.busRequest
                            targetServerIdx :: (Maybe (MkServerIdx numServers)) = serverMap targetAddr
                        case targetServerIdx of
                            Just serverIdx -> do
                                let
                                    targetServerArbiter :: Arbiter.Arbiter_IFC numClients
                                    targetServerArbiter = (select requestArbiterByServer serverIdx)
                                    arbiterClientSlot :: Arbiter.ArbiterClient_IFC
                                    arbiterClientSlot = (select targetServerArbiter.clients clientIdx)
                                arbiterClientSlot.request
                            Nothing -> do
                                let
                                    targetClientResponseArbiter :: Arbiter.Arbiter_IFC (TAdd numServers 1)
                                    targetClientResponseArbiter =  (select responseArbiterByClient clientIdx)

                                    clientResponseArbiterSlot :: Arbiter.ArbiterClient_IFC
                                    -- arbiters 0 to n-1 where `n:=numServer` are reserved
                                    -- for servers to make requests to. Arbiter n is reserved for
                                    -- when this rule needs to skip making a request to a server
                                    -- and should instead forward the `BusError UnMapped` response
                                    -- back to the client. Vector.last selects arbiter `n`
                                    clientResponseArbiterSlot = Vector.last targetClientResponseArbiter.clients
                                let
                                    responseUnMapped = case clientRequest.busRequest of
                                        BusReadRequest _ -> BusReadResponse (Left UnMapped)
                                        BusWriteRequest _ -> BusWriteResponse (Left UnMapped)
                                    response :: TaggedBusResponse inFlightTransactions
                                    response = TaggedBusResponse {
                                        tag = clientRequest.tag;
                                        busResponse =  responseUnMapped
                                        }
                                clientResponseArbiterSlot.request

                    (sprintf "client[%d] arbit submission" clientIdx): when True ==> do
                        let
                            selectedClientResponseArbiter :: Arbiter.Arbiter_IFC (TAdd numServers 1)
                            selectedClientResponseArbiter = (select responseArbiterByClient clientIdx)

                            selectedClientResponseQueue :: FIFOF (TaggedBusResponse inFlightTransactions)
                            selectedClientResponseQueue = (select clientResponseQueues clientIdx)

                            -- `TAdd numServers 1` because we can receive request from all servers
                            -- as well as a bypass requests from our one corresponding client request
                            -- queue
                            grantedIdx :: UInt (TLog (TAdd numServers 1))
                            grantedIdx = unpack selectedClientResponseArbiter.grant_id

                            isClientRequest :: Bool
                            isClientRequest = grantedIdx == fromInteger (valueOf numServers)
                        if isClientRequest then do
                            let
                                clientRequest :: TaggedBusRequest inFlightTransactions
                                clientRequest = selectedClientRequestQueue.first

                                responseUnMapped :: BusResponse
                                responseUnMapped = case clientRequest.busRequest of
                                    BusReadRequest _ -> BusReadResponse (Left UnMapped)
                                    BusWriteRequest _ -> BusWriteResponse (Left UnMapped)

                                response :: TaggedBusResponse inFlightTransactions
                                response = TaggedBusResponse {
                                    tag = clientRequest.tag;
                                    busResponse = responseUnMapped
                                }
                            selectedClientResponseQueue.enq response
                            selectedClientRequestQueue.deq
                        else do
                            let
                                grantedServerIdx :: MkServerIdx numServers
                                grantedServerIdx = truncate grantedIdx

                                selectedSubmitResponseQueue :: FIFOF (
                                                                    MkClientTagType numClients,
                                                                    TaggedBusResponse inFlightTransactions
                                                                )
                                selectedSubmitResponseQueue = (select submitResponseQueues grantedServerIdx)

                                response :: (MkClientTagType numClients, TaggedBusResponse inFlightTransactions)
                                response = selectedSubmitResponseQueue.first
                            selectedClientResponseQueue.enq response.snd
                            selectedSubmitResponseQueue.deq

    let serverRules :: Vector numServers (Rules)
        serverRules = genWith $ \serverIdx ->
            let
                selectedServerArbiter :: Arbiter.Arbiter_IFC numClients
                selectedServerArbiter = (select requestArbiterByServer serverIdx)

                selectedConsumeRequestQueue :: FIFOF (
                                                MkClientTagType numClients,
                                                TaggedBusRequest inFlightTransactions
                                                )
                selectedConsumeRequestQueue = (select consumeRequestQueues serverIdx)
            in
                rules
                    (sprintf "server[%d] arbit requests" serverIdx): when True ==> do
                        let
                            grantedClientIdx :: MkClientTagType numClients
                            grantedClientIdx = unpack selectedServerArbiter.grant_id

                            selectedClientRequestQueue :: FIFOF (TaggedBusRequest inFlightTransactions)
                            selectedClientRequestQueue = (select clientRequestQueues grantedClientIdx)

                            clientRequest :: TaggedBusRequest inFlightTransactions
                            clientRequest = selectedClientRequestQueue.first

                        selectedConsumeRequestQueue.enq (grantedClientIdx, clientRequest)
                        selectedClientRequestQueue.deq

                    (sprintf "server[%d] route response" serverIdx): when True ==> do
                        let
                            selectedSubmitResponseQueue :: FIFOF (
                                                                MkClientTagType numClients,
                                                                TaggedBusResponse inFlightTransactions
                                                            )
                            selectedSubmitResponseQueue = (select submitResponseQueues serverIdx)

                            response :: (MkClientTagType numClients, TaggedBusResponse inFlightTransactions)
                            response = selectedSubmitResponseQueue.first

                            clientTag :: MkClientTagType numClients
                            clientTag = response.fst

                            targetClientResponseArbiter :: Arbiter.Arbiter_IFC (TAdd numServers 1)
                            targetClientResponseArbiter = (select responseArbiterByClient clientTag)

                            arbiterClientSlot :: Arbiter.ArbiterClient_IFC
                            arbiterClientSlot = (select targetClientResponseArbiter.clients serverIdx)
                        arbiterClientSlot.request

    addRules |> foldr (<+>) (rules {}) clientRules
    addRules |> foldr (<+>) (rules {}) serverRules

    -- Client interface vector
    let clients :: Vector numClients (BusClient inFlightTransactions)
        clients = genWith $ \clientIdx ->
            let
                selectedClientRequestQueue :: FIFOF (TaggedBusRequest inFlightTransactions)
                selectedClientRequestQueue = (select clientRequestQueues clientIdx)

                selectedClientResponseQueue :: FIFOF (TaggedBusResponse inFlightTransactions)
                selectedClientResponseQueue = (select clientResponseQueues clientIdx)

                selectedTagEngine :: TagEngine inFlightTransactions
                selectedTagEngine = (select tagEngineByClientVec clientIdx)
            in
                interface BusClient
                    submitRequest   :: BusRequest
                                    -> ActionValue (MkTagType inFlightTransactions)
                    submitRequest busRequest = do
                        tag <- selectedTagEngine.requestTag
                        let taggedReuqest = TaggedBusRequest {tag = tag; busRequest = busRequest}
                        selectedClientRequestQueue.enq taggedReuqest
                        return tag

                    consumeResponse :: ActionValue (TaggedBusResponse inFlightTransactions)
                    consumeResponse = do
                        let
                            busResponse :: (TaggedBusResponse inFlightTransactions)
                            busResponse = selectedClientResponseQueue.first
                        selectedTagEngine.retireTag busResponse.tag
                        selectedClientResponseQueue.deq
                        return busResponse

    -- Server interface vector
    let servers :: Vector numServers (BusServer inFlightTransactions numClients)
        servers = genWith $ \serverIdx ->
            let
                selectedConsumeRequestQueue :: FIFOF (
                                                MkClientTagType numClients,
                                                TaggedBusRequest inFlightTransactions
                                                )
                selectedConsumeRequestQueue = (select consumeRequestQueues serverIdx)
            in
                interface BusServer
                    consumeRequest  :: ActionValue (
                                        MkClientTagType numClients,
                                        TaggedBusRequest inFlightTransactions
                                        )
                    consumeRequest = do
                        selectedConsumeRequestQueue.deq
                        return selectedConsumeRequestQueue.first

                    submitResponse  :: ( MkClientTagType numClients,
                                        TaggedBusResponse inFlightTransactions
                                        ) -> Action
                    submitResponse (clientTag, taggedBusResponse) = do
                        let
                            selectedSubmitResponseQueue :: FIFOF (
                                                                MkClientTagType numClients,
                                                                TaggedBusResponse inFlightTransactions
                                                            )
                            selectedSubmitResponseQueue = (select submitResponseQueues serverIdx)
                        selectedSubmitResponseQueue.enq (clientTag, taggedBusResponse)

    return $
        interface Bus
            clients = clients
            servers = servers