/*******************************************************************************
 * Copyright (C) 2024 Juan Neyra
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 ******************************************************************************/
#include "NetworkInterfaceTlm.h"
#include <iostream>

#include "ratatoskrUtils/traffic/Flit.h"
#include "utils/configuration.h"
#include "utils/utils.h"

using namespace std;

NetworkInterfaceTlm::NetworkInterfaceTlm(sc_module_name nm, Node& node,
        uint8_t max_pos[3]) : NetworkInterface(nm, node),
            credit_counter(NUM_CREDITS), ni_name(nm),
            init_peq(this, &NetworkInterfaceTlm::init_peq_cb),
            target_peq(this, &NetworkInterfaceTlm::target_peq_cb),
            curr_req(0), resp_in_progress(false),
            nxt_resp_pend(0) {
    sc_report_handler::set_actions(NI_LOG, SC_INFO, SC_LOG|SC_DISPLAY);
    try {
        this->id = node.id%(globalResources.nodes.size()/2);
        this->dbid = rep.registerElement("ProcessingElement", this->id);
        this->node = node;
        this->packetPortContainer = new PacketPortContainer(
                ("NI_PACKET_CONTAINER"+to_string(this->id)).c_str());
        copy(max_pos, max_pos + 3, this->max_pos);
    }
    catch (exception& e) {
        log_error(e.what());
    }
    SC_THREAD(thread);
    sensitive << clk.pos() << clk.neg();

    SC_METHOD(receivePacketFromPE);
    sensitive << packetPortContainer->portValidIn.pos();

    for(uint8_t lay=0; lay<NUM_ACC_LAYERS; lay++){
        initiator[lay] = new ni_init_socket((string(nm)+"_"+
                            to_string(lay)+"_init").c_str());
        initiator[lay]->register_nb_transport_bw(this, 
                        &NetworkInterfaceTlm::nb_transport_bw_cb, lay);

        target[lay] = new ni_targ_socket((string(nm)+"_"+
                            to_string(lay)+"targ").c_str());
        target[lay]->register_nb_transport_fw(this, 
                        &NetworkInterfaceTlm::nb_transport_fw_cb, lay);
    }
}

NetworkInterfaceTlm::~NetworkInterfaceTlm() {
    delete packetPortContainer;
}

void NetworkInterfaceTlm::initialize(){
    /* NOT IMPLEMENTED */
}


void NetworkInterfaceTlm::bind(Connection* con, SignalContainer* sigContIn, SignalContainer* sigContOut) {
    if (con==nullptr)
        packetPortContainer->bind(sigContIn, sigContOut);
}

void NetworkInterfaceTlm::receiveFlitFromRouter(){
    /* NOT IMPLEMENTED */
}

void NetworkInterfaceTlm::receivePacketFromPE() {
    if (packetPortContainer->portValidIn.posedge()) {
        //log_info("receive()");
        PacketCS* p = dynamic_cast<PacketCS*>(
                            packetPortContainer->portDataIn.read());
        generateFlitsForPacket(p);
        packet_send_queue.push(p);
    }
}

void NetworkInterfaceTlm::get_start_end_pos_string(PacketCS* p,
                                                string pos[2]){
    pos[0] = to_string(node.pos.x) + "," +
            to_string(node.pos.y) + "," + to_string(node.pos.z);
    pos[1] = to_string(p->dst.pos.x) + "," +
            to_string(p->dst.pos.y) + "," +
            to_string(p->dst.pos.z);
}

void NetworkInterfaceTlm::generateFlitsForPacket(Packet* p){
    /* NOT IMPLEMENTED */
}

void NetworkInterfaceTlm::generateFlitsForPacket(PacketCS* p) {
    int flitsPerPacket = p->size;
    for (int i = 0; i<flitsPerPacket; ++i) {
        FlitType flitType;
        if (flitsPerPacket==1)
            flitType = SINGLE;
        else if (i%flitsPerPacket==0)
            flitType = HEAD;
        else if (i%flitsPerPacket==flitsPerPacket-1)
            flitType = TAIL;
        else
            flitType = BODY;
        int seqNum = i%flitsPerPacket;
        Flit* current_flit = new Flit(flitType, seqNum, p, p->dataType, sc_time_stamp().to_double());
        p->toTransmit.push_back(current_flit->id);
        p->flits.push_back(current_flit);
        
        string start_end_pos[2];
        get_start_end_pos_string(p, start_end_pos);
        log_info("Flits generated for packet of id "+to_string(p->id)+
         " to be sent to "+start_end_pos[1]+" from "+start_end_pos[0]);
    }
}

tlm_gp* NetworkInterfaceTlm::build_transaction(PacketCS* p, Flit* f){
    // allocate tlm object
    tlm_gp* trans = m_mm.allocate();
    trans->acquire();
    // add id to transaction
    link_extension* ext = new link_extension();
    ext->link = Direction::local;
    ext->data_type = p->dataType;
    ext->packet_id = p->id;
    ext->flit_id = f->id;
    trans->set_extension(ext);

    trans->set_command(tlm::TLM_WRITE_COMMAND);

    float dst_pos[3] = {p->dst.pos.x, p->dst.pos.y, p->dst.pos.z};
    uint8_t int_dest_pos[3];
    convert_pos_to_int(dst_pos, int_dest_pos);
    trans->set_address(int_dest_pos[0] + int_dest_pos[1]*max_pos[0] +
                       int_dest_pos[2]*max_pos[0]*max_pos[1]);
    trans->set_data_ptr(reinterpret_cast<unsigned char*>(f));
    trans->set_data_length(p->size);
    trans->set_streaming_width(4);
    trans->set_byte_enable_ptr(0);
    trans->set_dmi_allowed(false);
    trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);
    curr_req = trans;

    return trans;
}


void NetworkInterfaceTlm::send_flit(PacketCS* p, Flit* f){
    tlm_gp* trans = build_transaction(p, f);

    sc_time delay = sc_time(REQ_INIT_DELAY, UNITS_DELAY);
    tlm::tlm_phase send_phase = tlm::BEGIN_REQ;

    string type_name = get_type_name(p->dataType);
    int lay = type_name == TYPE_STREAM ? 1:0;
    log_info("Sending flit of type "+ type_name + " to layer "+to_string(lay));
    (*initiator[lay])->nb_transport_fw(*trans, send_phase, delay);
    update_credits(*trans, -1);
}


int NetworkInterfaceTlm::get_type_from_extension(tlm_gp& trans){
    link_extension* extension;
    trans.get_extension<link_extension>(extension);
    return extension->data_type;
}


int NetworkInterfaceTlm::get_packet_id_from_extension(tlm_gp& trans){
    link_extension* extension;
    trans.get_extension<link_extension>(extension);
    return extension->packet_id;
}

bool NetworkInterfaceTlm::check_cs_needed(tlm_gp& trans){
    int type = get_type_from_extension(trans);
    return get_type_name(type) == TYPE_ROUT_CONFIG;
}

void NetworkInterfaceTlm::send_cs_rout_conf_msg(tlm_gp& trans){
    tlm_gp* new_trans = m_mm.allocate();
    // TODO: maybe use deep_copy_from function
    new_trans->acquire();
    new_trans->set_command(tlm::TLM_WRITE_COMMAND);
    new_trans->set_address(trans.get_address());
    new_trans->set_data_ptr(trans.get_data_ptr());
    new_trans->set_data_length(trans.get_data_length());
    new_trans->set_streaming_width(trans.get_streaming_width());
    new_trans->set_byte_enable_ptr(trans.get_byte_enable_ptr());
    new_trans->set_dmi_allowed(false);
    new_trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);

    // add id to new_transaction
    link_extension* ext = new link_extension();
    ext->link = Dir::local; // set direction of connected router
    ext->data_type = get_type_from_extension(trans);
    ext->packet_id = -1;
    ext->flit_id = -1;
    new_trans->set_extension(ext);

    // send transaction in socket
    tlm_phase phase = BEGIN_REQ;
    sc_time delay = sc_time(REQ_INIT_DELAY, UNITS_DELAY);
    (*initiator[1])->nb_transport_fw(*new_trans, phase, delay);
}

void NetworkInterfaceTlm::send_data_to_noc(){
    // if no flits to transmit, generate them
    PacketCS* p = packet_send_queue.front();
    if (p->toTransmit.empty()){
        generateFlitsForPacket(p);
    }
    // get current flit to send
    flitID_t f_id = p->toTransmit.front();
    auto iter = find_if(p->flits.begin(), p->flits.end(),
            [&f_id](Flit* f) { return f->id==f_id; });
    Flit* current_flit = *iter;
    current_flit->injectionTime = sc_time_stamp().to_double();

    auto toDelete_pos = find(p->toTransmit.begin(), p->toTransmit.end(), current_flit->id);
    p->toTransmit.erase(toDelete_pos);
    p->inTransmit.push_back(current_flit->id);

    if (p->toTransmit.empty()) {
        packet_send_queue.pop();
    }

    send_flit(p, current_flit);

    string start_end_pos[2];
    get_start_end_pos_string(p, start_end_pos);
    log_info("Flit sent to "+start_end_pos[1]+" from "+start_end_pos[0]);
    globalReport.issueNoCInputDataAmount(sc_time_stamp(),globalResources.bitWidth);
}


void NetworkInterfaceTlm::thread() {
    bool log_wait = true;
    while(true){
        //log_info("send_data_process()");
        if (clk.posedge()) {
            if (!packet_send_queue.empty()) {
                if (credit_counter != 0 && curr_req == 0){
                    send_data_to_noc();
                    log_wait = true;
                }
                else {
                    if(log_wait){
                        if(credit_counter == 0){
                            log_info("No more credits. Waiting for Router!");
                        }
                        else{
                            log_info("Transaction in progress. Waiting for Router!");
                        }
                        log_wait = false;
                    }
                }
            }
            if (!packet_recv_queue.empty()) {
                if (packetPortContainer->portFlowControlIn.read()) {
                    PacketCS* p = packet_recv_queue.front();
                    packet_recv_queue.pop();
                    packetPortContainer->portValidOut.write(true);
                    packetPortContainer->portDataOut.write(p);
                    // log arrival of packet
                    string start_end_pos[2];
                    get_start_end_pos_string(p, start_end_pos);
                    log_info("Packet with id " + to_string(p->id) + 
                            " arrived at PE "+start_end_pos[1]);
                }
            }
        }
        else if (clk.negedge()) {
            packetPortContainer->portValidOut.write(false);
        }
        wait();
    }
}

tlm::tlm_sync_enum NetworkInterfaceTlm::nb_transport_bw_cb(int id,
                            tlm::tlm_generic_payload& trans,
                            tlm::tlm_phase& phase, sc_time& delay) {
    int p_id = get_packet_id_from_extension(trans);
    log_info("Backward transport callback start from layer " + 
            to_string(id) + ". Packet id: " + to_string(p_id));
    log_info("Sanity check: Phase " + string(phase.get_name()) + ", "
            + "Delay: " + delay.to_string());
    init_peq.notify(trans, phase, delay);
    return tlm::TLM_ACCEPTED;
}


void NetworkInterfaceTlm::init_peq_cb(tlm_gp& trans, const tlm::tlm_phase& phase){
    int p_id = get_packet_id_from_extension(trans);
    log_info("Init PEQ callback start for packet "+to_string(p_id));
    log_info("Phase "+string(phase.get_name())+" received");
    if (phase == tlm::END_REQ) {
        update_credits(trans, 1);
        curr_req = 0;
        if(p_id == -1){
            // Config msg transaction ended, so packet can be scheduled for deletion
            trans.set_response_status(tlm::TLM_OK_RESPONSE);
            unsigned char* data_ptr = trans.get_data_ptr();
            Flit* received_flit = reinterpret_cast<Flit*>(data_ptr);
            PacketCS* p = dynamic_cast<PacketCS*>(received_flit->packet);
            packet_recv_queue.push(p);
        }
    }
    else if (phase == tlm::BEGIN_REQ || phase == tlm::END_RESP)
        log_fatal("Illegal transaction phase received by initiator");

    if (phase == tlm::BEGIN_RESP) {
        check_transaction(trans);
        // Send final phase transition to target
        tlm::tlm_phase fw_phase = tlm::END_RESP;
        sc_time delay = sc_time(RESP_END_DELAY, UNITS_DELAY);
        (*initiator[0])->nb_transport_fw(trans, fw_phase, delay);
    }
}


// Called on receiving BEGIN_RESP or TLM_COMPLETED
void NetworkInterfaceTlm::check_transaction(tlm_gp& trans) {
    if(trans.is_response_error()) {
        char txt[100];
        sprintf(txt, "Transaction returned with error, response status = %s",
                    trans.get_response_string().c_str());
        log_error(txt);
    }

    // Log completed routing
    tlm_command cmd = trans.get_command();
    sc_dt::uint64    adr = trans.get_address();
    int*             ptr = reinterpret_cast<int*>(trans.get_data_ptr());
    std::stringstream stream;
    stream << hex << adr << " check, cmd=" << (cmd ? 'W' : 'R')
                         << ", data=" << hex << *ptr;
    log_info(stream.str());
    // Allow the memory manager to free the transaction object
    trans.release();
}

tlm::tlm_sync_enum NetworkInterfaceTlm::nb_transport_fw_cb(int id,
            tlm_gp& trans, tlm::tlm_phase& phase, sc_time& delay){
    int p_id = get_packet_id_from_extension(trans);
    log_info("Forward transport callback start from layer " + 
            to_string(id) + ". Packet id: " + to_string(p_id));
    log_info("Sanity check: Phase " + string(phase.get_name()) + ", "
            + "Delay: " + delay.to_string());
    unsigned int     len = trans.get_data_length();
    unsigned char*   byt = trans.get_byte_enable_ptr();
    unsigned int     wid = trans.get_streaming_width();
    if (byt != 0) {
        trans.set_response_status( tlm::TLM_BYTE_ENABLE_ERROR_RESPONSE );
        return tlm::TLM_COMPLETED;
    }
    if (len > 4 || wid < len) {
        trans.set_response_status( tlm::TLM_BURST_ERROR_RESPONSE );
        return tlm::TLM_COMPLETED;
    }

    target_peq.notify(trans, phase, delay);
    return tlm::TLM_ACCEPTED;
}


void NetworkInterfaceTlm::receive_and_process_flit(tlm_gp& trans){
    trans.set_response_status(tlm::TLM_OK_RESPONSE);
    unsigned char* data_ptr = trans.get_data_ptr();
    Flit* received_flit = reinterpret_cast<Flit*>(data_ptr);
    PacketCS* p = dynamic_cast<PacketCS*>(received_flit->packet);
    double time = sc_time_stamp().to_double();
    auto position = find(p->inTransmit.begin(), p->inTransmit.end(), received_flit->id);
    if (position!=p->inTransmit.end())
        p->inTransmit.erase(position);
    p->transmitted.push_back(received_flit->id);

    // log arrived flit
    if (received_flit->type==TAIL || received_flit->type==SINGLE){
        stringstream ss;
        string flit_type = received_flit->type==SINGLE ? "Single":"Tail";
        ss << "Receive Flit " << flit_type << *received_flit;
        log_info(ss.str());
    }
    if (!p->toTransmit.empty() || !p->inTransmit.empty()) {
        stringstream ss;
        ss << "Received Tail Flit, but still missing flits! "
                << *received_flit;
        log_info(ss.str());
    }
    //globalReport.issueNoCOutputDataAmount(sc_time_stamp(),globalResources.bitWidth);
    if (p->toTransmit.empty() && p->inTransmit.empty()){
        // send config message when all flits of packet arrive
        if(check_cs_needed(trans)){
            log_info("Send router configuration message. Packet id: "
                        + to_string(p->id));
            send_cs_rout_conf_msg(trans);
        }
        else{
            packet_recv_queue.push(p);
        }
    }
    //credit_counter++; no need to update credit in target (unless receiving response)
    //log_info("Credit counter updated: "+to_string(credit_counter));
    
    // send response if not type stream
    if(get_type_name(p->dataType) != TYPE_STREAM){
        if(resp_in_progress) {
            if(nxt_resp_pend){
                log_fatal("Attempt to have two pending responses in target");
            }
            log_info("Previous response transaction still in progress, scheduling it for next time");
            nxt_resp_pend = &trans;
        }
        else{ send_response(trans); }
    }
}


void NetworkInterfaceTlm::target_peq_cb(tlm_gp& trans, const tlm::tlm_phase& phase){
    int p_id = get_packet_id_from_extension(trans);
    log_info("Target PEQ callback start for packet "+to_string(p_id));
    log_info("Phase "+string(phase.get_name())+" received");
    switch (phase) {
        case tlm::BEGIN_REQ:
            trans.acquire();
            send_end_req(trans);
            break;

        case tlm::END_RESP:
            if (!resp_in_progress){
                log_fatal(
                "Illegal transaction phase END_RESP received by target");
            }
            trans.release();
            update_credits(trans, 1);
            // Target itself is now clear to issue the next BEGIN_RESP
            resp_in_progress = false;
            if (nxt_resp_pend){
                send_response(*nxt_resp_pend);
                nxt_resp_pend = 0;
            }
            break;

        case tlm::END_REQ:
        case tlm::BEGIN_RESP:
            log_fatal("Illegal transaction phase received by target");
            break;

        default:
            if(phase == INTERNAL_PROC_PHASE){
                receive_and_process_flit(trans);
            }
            break;
    }
}

tlm::tlm_sync_enum NetworkInterfaceTlm::send_end_req(tlm_gp& trans){
    // Queue the acceptance and the response with the appropriate latency
    tlm::tlm_phase bw_phase = tlm::END_REQ;
    sc_time delay = sc_time(REQ_END_DELAY, UNITS_DELAY); // Accept delay

    tlm::tlm_sync_enum status = (*target[0])->nb_transport_bw(trans, 
                                                bw_phase, delay);
    if (status == tlm::TLM_COMPLETED) {
        trans.release();
        curr_req = 0;
        return status;
    }

    // Queue internal event to mark beginning of response
    int type = get_type_from_extension(trans);
    delay = get_type_name(type) != TYPE_STREAM ? 
                delay + sc_time(INTERN_PROC_DELAY, UNITS_DELAY) : 
                SC_ZERO_TIME; // no processing on stream
    target_peq.notify(trans, INTERNAL_PROC_PHASE, delay);

    return status;
}


void NetworkInterfaceTlm::send_response(tlm_gp& trans){
    int p_id = get_packet_id_from_extension(trans);
    log_info("Send response for packet "+to_string(p_id));
    resp_in_progress = true;
    sc_time delay = SC_ZERO_TIME;
    tlm::tlm_phase bw_phase = tlm::BEGIN_RESP;

    tlm::tlm_sync_enum status = (*target[0])->nb_transport_bw(
                                    trans, bw_phase, delay);
    update_credits(trans, -1);

    if (status == tlm::TLM_UPDATED){
      target_peq.notify(trans, bw_phase, delay);
    }
    else if (status == tlm::TLM_COMPLETED) {
      // The initiator has terminated the transaction
      trans.release();
      resp_in_progress = false;
      update_credits(trans, 1);
      curr_req = 0;
    }
}

void NetworkInterfaceTlm::update_credits(tlm_gp&trans, int added_creds){
    string type_name = get_type_name(get_type_from_extension(trans));
    int p_id = get_packet_id_from_extension(trans);

    if(type_name == TYPE_PACKET || (type_name == TYPE_ROUT_CONFIG && p_id != -1)){
        credit_counter += added_creds;
        log_info("Credit counter updated: "+to_string(credit_counter));
    }
}

void NetworkInterfaceTlm::log_info(string msg){
    string log_msg = ni_name + ": (Node" + 
                        to_string(node.id)+"): "+msg;
    SC_REPORT_INFO(NI_LOG, (log_msg).c_str());
}

void NetworkInterfaceTlm::log_error(string msg){
    string log_msg = ni_name + ": (Node" + 
                        to_string(node.id)+"): "+msg;
    SC_REPORT_ERROR(NI_LOG, (log_msg).c_str());
}

void NetworkInterfaceTlm::log_fatal(string msg){
    string log_msg = ni_name + ": (Node" + 
                        to_string(node.id)+"): "+msg;
    SC_REPORT_FATAL(NI_LOG, (log_msg).c_str());
}