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

import Types
import BusTypes
import Client
import Server
import TagEngine

import Vector
import Util
import Arbiter
import FIFO
import FIFOF
import SpecialFIFOs
import Printf

-- Creates a Bus Module that supports multiple clients and servers
-- submitting requests and simultaneously returning responses.
-- Responses can be consumed by clients out of order(useful when some
-- servers respond before others) as all client submitted requests are
-- tagged with tags that are unique for the duration of the transaction
--  - and well-behaved servers should keep that tag when responding.

-- Explicitly inform the compiler that log2 n <= log2(n + 1)
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 client
    -- response arbiter as there are `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 submitter to
    -- the client response arbiter as it may determine that a request is
    -- unmappable and simply bypass submitting the request to a server,
    -- instead opting to form a `BusError UnMapped` response to be submitted
    -- directly to a client response arbiter. Thus the client response arbiter
    -- 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)

    -- 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 single depth FIFOs make it easier to push/pull data
    -- to/from internal server methods as they provide back-pressure in both directions,
    -- and behave as a wire when queue is empty.
    -- If looking at the example bus.drawio diagram, the following two vectors effectively
    -- correspond to the two arrows going from the blue box to the servers.
    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 selectedClientReqQueue :: FIFOF (TaggedBusRequest inFlightTransactions)
                selectedClientReqQueue = select clientRequestQueues clientIdx

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

                selectedClientRespArbiter :: Arbiter.Arbiter_IFC (TAdd numServers 1)
                selectedClientRespArbiter = select responseArbiterByClient clientIdx

                clientRouterRule :: Rules
                clientRouterRule =  clientRouteRequest
                    clientIdx
                    selectedClientReqQueue
                    requestArbiterByServer
                    responseArbiterByClient
                    serverMap

                clientArbiterRule :: Rules
                clientArbiterRule = clientArbitSubmission
                    clientIdx
                    selectedClientReqQueue
                    selectedClientRespQueue
                    selectedClientRespArbiter
                    submitResponseQueues
            in
                clientRouterRule <+> clientArbiterRule

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

                selectedConsumeReqQueue :: FIFOF (MkClientTagType numClients, TaggedBusRequest inFlightTransactions)
                selectedConsumeReqQueue = select consumeRequestQueues serverIdx

                selectedSubmitRespQueue :: FIFOF (MkClientTagType numClients, TaggedBusResponse inFlightTransactions)
                selectedSubmitRespQueue = select submitResponseQueues serverIdx

                serverRouterRule :: Rules
                serverRouterRule = serverRouteResponse
                    serverIdx
                    selectedSubmitRespQueue
                    responseArbiterByClient

                serverArbiterRule :: Rules
                serverArbiterRule = serverArbitRequests
                    serverIdx
                    selectedServerArbiter
                    selectedConsumeReqQueue
                    clientRequestQueues
            in
                serverRouterRule <+> serverArbiterRule

    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