#include "Sensors.h"
#include <Adafruit_EEPROM_I2C.h>

#define RFM95_CS    8
#define RFM95_INT   3
#define RFM95_RST   4
#define RF95_FREQ 915.0

#define TEMPORARY_ID 100
#define MASTER_NODE_ID 0
#define BROADCAST_ADDRESS 255

#define MAX_NEIGHBORS 15
#define MAX_HOPS 20

#define MESSAGE_TYPE_BROADCAST 0
#define MESSAGE_TYPE_COMMAND 1
#define MESSAGE_TYPE_SENSOR_DATA 2
#define MESSAGE_TYPE_ACK 4
#define MESSAGE_TYPE_NETWORK_ADDITION_PROPOSAL 10
#define MESSAGE_TYPE_NETWORK_NEIGHBOUR_UPDATE 11
#define MESSAGE_TYPE_NETWORK_ROUTE_REQUEST 12


#define MESSAGE_CONSUMED 5
#define MESSAGE_NOT_CONSUMED 6

Adafruit_EEPROM_I2C EEPROM;

#define EEPROM_ADDR 0x50  

RH_RF95 rf95(RFM95_CS, RFM95_INT);
uint8_t NODE_ID;
unsigned long lastHandshakeAttempt = 0;
bool inHandshakeProcess = false;
RHReliableDatagram manager(rf95, NODE_ID);

Neighbor neighbors[MAX_NEIGHBORS];
uint8_t neighborCount = 0;
uint8_t** masterPaths;
 
Sensors sensorModule;


uint8_t** getPath(uint8_t nodeID);

void setup() {
    Serial.begin(4800);
    delay(10);
    Serial.println("LoRa NODE");


    if (EEPROM.begin(EEPROM_ADDR)) {  // you can stick the new i2c addr in here, e.g. begin(0x51);
        Serial.println("Found I2C EEPROM");
    } else {
        Serial.println("I2C EEPROM not identified ... check your connections?\r\n");
        while (1) delay(10);
    }


    masterPaths = (uint8_t**) malloc(2 * sizeof(uint8_t*));
    for (int i = 0; i < 2; i++) {
        masterPaths[i] = (uint8_t*) malloc(10 + (i * 10) * sizeof(uint8_t));
        for(int x = 0; x < 10 + (i * 10); x++){
            masterPaths[i][x] = 100;
        }
    } 

    pinMode(RFM95_RST, OUTPUT);
    digitalWrite(RFM95_RST, HIGH);

    Serial.println("Arduino LoRa TX!");

    digitalWrite(RFM95_RST, LOW);
    delay(10);
    digitalWrite(RFM95_RST, HIGH);
    delay(10);

    while (!rf95.init()) {
        while (1);
    }
    // Defaults after init are 434.0MHz, modulation GFSK_Rb250Fd250, +13dbM
    if (!rf95.setFrequency(RF95_FREQ)) {
        while (1);
    }
    // Defaults after init are 434.0MHz, 13dBm, Bw = 125 kHz, Cr = 4/5, Sf = 128chips/symbol, CRC on
    // The default transmitter power is 13dBm, using PA_BOOST.
    // If you are using RFM95/96/97/98 modules which uses the PA_BOOST transmitter pin, then
    // you can set transmitter powers from 5 to 23 dBm:
    rf95.setTxPower(20, false);
        // while (1);
    sensorModule.initialize(); 
  
    uint8_t eepromID = EEPROM.read(0);
  
    if(eepromID > 0 && eepromID != TEMPORARY_ID) {
        NODE_ID = eepromID;
        manager.setThisAddress(NODE_ID);
    } else {
        // Initiate handshake process
        NODE_ID = TEMPORARY_ID;
        manager.setThisAddress(NODE_ID);
        inHandshakeProcess = true;
    }
    sensorModule.setAddress(NODE_ID);
    
}

void loop() {
  // Handle incoming messages
    if (manager.available()) {
        Message msg;
        uint8_t len = sizeof(Message);
        uint8_t from;
        if (manager.recvfrom((uint8_t*)&msg, &len, &from)) {
            Serial.println("Msg Recv");
            handleReceivedMessage(msg);
        }
    }

    // Broadcast presence

    static signed long lastBroadcast = -1800000; // 30 mins
        Serial.print(accurateMillis());
        Serial.print(" ");
        Serial.println(lastBroadcast);
    if (accurateMillis() - lastBroadcast >= 1800000) { 
        Serial.println("BroadCasting Presence");
        broadcastPresence();
        lastBroadcast = accurateMillis();
    }
    // Check and prune neighbors list
    static signed long lastPrune = 0;
    if (accurateMillis() - lastPrune >= 7200000) { // 2 hours
        Serial.println("Prune Neighbours");
        pruneNeighbors();
        lastPrune = accurateMillis();
    }

    // send sensor report
    static signed long lastReport = -150000;
    if(inHandshakeProcess) {
        if(millis() - lastHandshakeAttempt > 150000) { // Every 2.5 minutes (handshake retries slower)
            Serial.println("Attempt Handshake");
            lastHandshakeAttempt = millis();
            attemptHandshake();
        }
    }else {
        if (accurateMillis() - lastReport >= (60000)) { // 1 min
            int messageCount = sensorModule.readSensorsAndSend();
            Serial.println("Broadcast Sensor Reading");
            for(int i = 0; i < messageCount; i++) {

                // Checking if any of the masterPaths are valid
                bool* pathValidity = isValidPath(masterPaths);
                Message msg = sensorModule.messages[i];
                if (pathValidity[0]) { // Shortest path is valid
                    for (int i = 0; i < 10; i++) {
                        msg.route[i] = masterPaths[0][i];
                    }
                    // Send the message to the first address in the route
                    manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[0]);
                } else if (pathValidity[1]) { // Reliable path is valid
                    for (int i = 0; i < 10; i++) {
                        msg.route[i] = masterPaths[1][i];
                    }
                    // Send the message to the first address in the route
                    manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[0]);
                }else {
                    sendToMasterOrBroadcast(msg);
                }
            }
            lastReport = accurateMillis();
        }
    }


    // Get Optimized Path to Master
    static signed long lastOptimized = -3600000 * 3;
    if (accurateMillis() - lastOptimized >= 3600000 * 3) { // 3 hours
        uint8_t** paths = getPath(MASTER_NODE_ID);

        if (isValidPath(paths)[0] || isValidPath(paths)[1]) { // if either shortest or reliable path present
            // Free old memory
            for (int i = 0; i < 2; i++) {
                free(masterPaths[i]);
            }
            free(masterPaths);
            
            // Assign new memory
            masterPaths = paths;
            lastOptimized = accurateMillis();
        }
    }
}

void handleReceivedMessage(Message msg) {
    //TODO: Add Id to list of last 10 commands, reject duplicates, trim list ever 5 minutes
    if(msg.senderID == NODE_ID){
        return;
    }
    Message ackMsg;
    strncpy(ackMsg.id, generateMessageID(), sizeof(ackMsg.id));
    ackMsg.type = MESSAGE_TYPE_ACK;
    ackMsg.targetID = msg.senderID;
    ackMsg.senderID = NODE_ID;
    ackMsg.hops = 0;

    if (msg.type == MESSAGE_TYPE_BROADCAST) { // Presence Acknolwedgment
        addNeighbor(msg.senderID);
        ackMsg.sensorValue = MESSAGE_CONSUMED;
        sendToTargetOrBroadcast(ackMsg);
        return;
    }

    if (msg.targetID == NODE_ID) { // Direct message
        if (msg.type == MESSAGE_TYPE_COMMAND) {
            // Check the data array when processing commands
            if (strncmp(msg.data, "CMD", 20) == 0) {
                // Execute the specific command for "CMD"
            }
            // Add more command checks as needed

            ackMsg.sensorValue = MESSAGE_CONSUMED;
            sprintf(ackMsg.data, "300", NULL);
        } else {
            ackMsg.sensorValue = MESSAGE_NOT_CONSUMED;
            sprintf(ackMsg.data, "500", NULL);
        }

        // Get the path using getPath and set the route for the acknowledgement message
        uint8_t** paths = getPath(msg.senderID);
        if (paths) {
            // Assuming we prefer the shortest path (paths[0]), but you can change this if needed
            for (int i = 0; i < 10; i++) { // Assuming route size is 10
                ackMsg.route[i] = paths[0][i];
            }

            // Freeing up the allocated memory
            for (int i = 0; i < 2; i++) {
                free(paths[i]);
            }
            free(paths);
        }

        // Send the message using the new route
        manager.sendto((uint8_t*)&ackMsg, sizeof(Message), ackMsg.route[0]);

        return;
    }
    // Increase hop count
    msg.hops++;
    if (msg.hops > MAX_HOPS) return; // Stop Infinite cycle

    if (msg.type == MESSAGE_TYPE_SENSOR_DATA) {
        msg.lastRelayID = NODE_ID;
        manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[msg.hops]);
        return;
    }

    relayMessage(msg);
}

void relayMessage(Message msg) {
    msg.lastRelayID = NODE_ID;
    if (msg.targetID != BROADCAST_ADDRESS) {

        // Send the message using pathed route
        manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[msg.hops]);
        const uint16_t ACK_TIMEOUT = 2500;  // 500 milliseconds
        // Wait for an acknowledgment with specific values
        unsigned long startTime = accurateMillis();
        while (accurateMillis() - startTime < ACK_TIMEOUT) {// 5 secs = 5000 milliseconds
            if (manager.available()) {
                // Read the incoming message
                Message incomingMsg;
                uint8_t len = sizeof(Message);
                uint8_t from;
                manager.recvfrom((uint8_t*)&incomingMsg, &len, &from);


                if (incomingMsg.senderID != msg.senderID && incomingMsg.targetID == NODE_ID && strcmp(incomingMsg.data, "300")) {
                    return;                        
                }
            }
        }
        
        // bool isClose = isNeighbor(msg.targetID);
        // const uint16_t ACK_TIMEOUT = 5000;  // 5000 milliseconds
        // uint64_t startTime = accurateMillis();
        // if(isClose){
        //     manager.sendto((uint8_t *)&msg, sizeof(Message), msg.targetID);

        //     unsigned long startTime = accurateMillis();
        //     while (accurateMillis() - startTime < ACK_TIMEOUT) {// 5 secs = 5000 milliseconds
        //         if (manager.available()) {
        //             // Read the incoming message
        //             Message incomingMsg;
        //             uint8_t len = sizeof(Message);
        //             uint8_t from;
        //             manager.recvfrom((uint8_t*)&incomingMsg, &len, &from);
        

        //             if (incomingMsg.type == MESSAGE_TYPE_ACK && incomingMsg.targetID == NODE_ID) {
        //                 return;                        
        //             }
        //         }
        //     }
        // }else {
        //     for (uint8_t i = 0; i < neighborCount; i++) {
        //         manager.sendto((uint8_t *)&msg, sizeof(Message), neighbors[i].nodeID);

        //         unsigned long startTime = accurateMillis();
        //         while (accurateMillis() - startTime < ACK_TIMEOUT) {// 5 secs = 5000 milliseconds
        //             if (manager.available()) {
        //                 // Read the incoming message
        //                 Message incomingMsg;
        //                 uint8_t len = sizeof(Message);
        //                 uint8_t from;
        //                 manager.recvfrom((uint8_t*)&incomingMsg, &len, &from);
        //                 if (incomingMsg.type == MESSAGE_TYPE_ACK && incomingMsg.targetID == NODE_ID) {
        //                     break;// continue for loop to send to neighbours                     
        //                 }
        //             }
        //         }
        //     }
        // }
    }
}

char* convertNeighborsToCharArray(Neighbor neighbors[], int count) {
    static char result[MAX_NEIGHBORS];
    for(int i = 0; i < count && i < MAX_NEIGHBORS; i++) {
        result[i] = (char)neighbors[i].nodeID;
    }
    return result;
}

void uploadNeighbours() {
    Message msg;
    
    // Fill message details
    msg.type = MESSAGE_TYPE_NETWORK_NEIGHBOUR_UPDATE;
    msg.targetID = MASTER_NODE_ID;

    strncpy(msg.id, generateMessageID(), sizeof(msg.id));
    msg.senderID = NODE_ID;
    // Get neighbors as char array
    char* neighborData = convertNeighborsToCharArray(neighbors, neighborCount);
    for(int i = 0; i < neighborCount && i < MAX_NEIGHBORS; i++) {
        msg.data[i] = neighborData[i];
    }

    // Assuming you have a function called sendToMasterOrBroadcast that takes a Message
    sendToMasterOrBroadcast(msg);
}


void broadcastPresence() {
    Message msg;
    msg.type = MESSAGE_TYPE_BROADCAST;
    msg.senderID = NODE_ID;
    msg.targetID = BROADCAST_ADDRESS;
    strncpy(msg.id, generateMessageID(), sizeof(msg.id));
    msg.hops = 0;
    // Fill other data fields if necessary
    
    manager.sendto((uint8_t*)&msg, sizeof(Message), BROADCAST_ADDRESS);
}

void addNeighbor(uint8_t id) {
    for (uint8_t i = 0; i < neighborCount; i++) {
        if (neighbors[i].nodeID == id) {
            neighbors[i].lastSeen = accurateMillis();
            return;
        }
    }

    if (neighborCount < MAX_NEIGHBORS) {
        neighbors[neighborCount].nodeID = id;
        neighbors[neighborCount].lastSeen = accurateMillis();
        neighborCount++;
    }
    uploadNeighbours();
}

void pruneNeighbors() {
    for (uint8_t i = 0; i < neighborCount; i++) {
        if (accurateMillis() - neighbors[i].lastSeen >= 3600000 / 4 * 3) { // 1 hours / 4 * 3 = 45 minsy
            for (uint8_t j = i; j < neighborCount - 1; j++) {
                neighbors[j] = neighbors[j + 1];
            }
            neighborCount--;
            i--;  // Check the same index again
        }
    }
    uploadNeighbours();
}
 

void sendToMasterOrBroadcast(Message& message) {
    if (isNeighbor(MASTER_NODE_ID)) {
        manager.sendto((uint8_t*)&message, sizeof(Message), MASTER_NODE_ID);
    } else {
        manager.sendto((uint8_t*)&message, sizeof(Message), BROADCAST_ADDRESS);
    }
}

void sendToTargetOrBroadcast(Message& message) {
    if (isNeighbor(message.targetID)) {
        manager.sendto((uint8_t*)&message, sizeof(Message), message.targetID);
    } else {
        manager.sendto((uint8_t*)&message, sizeof(Message), BROADCAST_ADDRESS);
    }
}
bool isNeighbor(uint8_t nodeId) {
    for (uint8_t i = 0; i < neighborCount; i++) {
        if (neighbors[i].nodeID == nodeId) {
            return true;
        }
    }
    return false;
}

uint64_t accurateMillis() {
    const uint64_t maxMillisValue = 0xFFFFFFFF; // max val of unsigned long
    const uint64_t overflowIntervalMillis = maxMillisValue + 1;
    static uint64_t lastMillisValue = 0;
    static short numOverflows = 0; // 7 overflows a year, 140 in 20 years, if this is still running jesus
    uint64_t currentMillisValue = millis();
    
    // rip overflowed
    if (currentMillisValue < lastMillisValue) {
        numOverflows++;
    }
    
    lastMillisValue = currentMillisValue;
    
    return (numOverflows * overflowIntervalMillis) + currentMillisValue;
}

void attemptHandshake() {
    uint8_t proposedID = generateID();
    Message msg;
    msg.type = MESSAGE_TYPE_NETWORK_ADDITION_PROPOSAL;
    msg.senderID = NODE_ID;
    msg.targetID = MASTER_NODE_ID;  
    snprintf(msg.data, sizeof(msg.data), "%u", proposedID);
    strncpy(msg.id, generateMessageID(), sizeof(msg.id));

    static signed long lastProposalTime = 0;
    if(isNeighbor(MASTER_NODE_ID)) {
        manager.sendto((uint8_t*)&msg, sizeof(Message), MASTER_NODE_ID);
    } else {
        manager.sendto((uint8_t*)&msg, sizeof(Message), BROADCAST_ADDRESS);
    }

    unsigned long startTime = accurateMillis();
    while (accurateMillis() - startTime < 60000) {// 1 minute = 60000 milliseconds
        if (manager.available()) {
            // Read the incoming message
            Message incomingMsg;
            uint8_t len = sizeof(Message);
            uint8_t from;
            manager.recvfrom((uint8_t*)&incomingMsg, &len, &from);
 

            if (incomingMsg.type == MESSAGE_TYPE_NETWORK_ADDITION_PROPOSAL && incomingMsg.targetID == NODE_ID) {
                if(incomingMsg.data[0] == '3' && incomingMsg.data[1] == '0' && incomingMsg.data[2] == '0') {
                    NODE_ID = proposedID;
                    EEPROM.write(0, NODE_ID);
                    manager.setThisAddress(NODE_ID);
                    sensorModule.setAddress(NODE_ID);
                    break;
                    // EEPROM.commit(); 
                    
                } else if(incomingMsg.data[0] == '5' && incomingMsg.data[1] == '0' && incomingMsg.data[2] == '0') {
                    break;
                    // Nothing specific to do here, a new ID will be proposed in the next attempt
                }
            }
        }
    }

}

uint8_t generateID() {
    uint8_t proposedID = random(1, 256);
    while(proposedID == TEMPORARY_ID || proposedID == 0) {
        proposedID = random(1, 256);
    }
    return proposedID;
}
char* generateMessageID() {
    static char id[10];  // 9 characters + null terminator
    const char* charset = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";

    for(int i = 0; i < 9; i++) {
        id[i] = charset[random(0, 62)];  // 62 possible alphanumeric characters
    }
    id[9] = '\0';  // Null terminate the string

    return id;
}


uint8_t** getPath(uint8_t nodeID) {
    Message msg;
    msg.type = MESSAGE_TYPE_NETWORK_ROUTE_REQUEST;
    msg.targetID = MASTER_NODE_ID;
    msg.senderID = NODE_ID;
    msg.sensorValue = (uint16_t)nodeID;  

    // Send the message

    // Checking if any of the masterPaths are valid
    bool* pathValidity = isValidPath(masterPaths);
    if (pathValidity[0]) { // Shortest path is valid
        for (int i = 0; i < 10; i++) {
            msg.route[i] = masterPaths[0][i];
        }
        // Send the message to the first address in the route
        manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[0]);
    } else if (pathValidity[1]) { // Reliable path is valid
        for (int i = 0; i < 10; i++) {
            msg.route[i] = masterPaths[1][i];
        }
        // Send the message to the first address in the route
        manager.sendto((uint8_t*)&msg, sizeof(Message), msg.route[0]);
    }else {
        sendToMasterOrBroadcast(msg);
    }

    free(pathValidity); // Free the validity array

    
 
    uint8_t** paths = (uint8_t**) malloc(2 * sizeof(uint8_t*));
    for (int i = 0; i < 2; i++) {
        paths[i] = (uint8_t*) malloc(10 + (i * 10) * sizeof(uint8_t));
        for(int x = 0; x < 10 + (i * 10); x++){
            paths[i][x] = 100;
        }
    }
 
    unsigned long startTime = millis();
    unsigned long timeout = 2000; // 2 seconds
    while (accurateMillis() - startTime < timeout) { 
        if (manager.available()) {
            // Read the incoming message
            Message receivedMsg;
            uint8_t len = sizeof(Message);
            uint8_t from;
            manager.recvfrom((uint8_t*)&receivedMsg, &len, &from);
 

            if (receivedMsg.type == MESSAGE_TYPE_NETWORK_ROUTE_REQUEST && receivedMsg.targetID == NODE_ID) {
                
                for (int i = 0; i < 10; i++) {
                    if(receivedMsg.route[i] == 100)
                        break;
                    paths[0][i] = (uint8_t)receivedMsg.route[i];
                }

                int dataIndex = 0;
                for (int i = 0; i < 20; i++) {
                    if(receivedMsg.data[i] == 100)
                        break;
                    paths[1][dataIndex++] = (uint8_t)receivedMsg.route[i];
                } 
                return paths;
            }
        } 
    }

    // Cleanup if no path received
    for (int i = 0; i < 2; i++) {
        free(paths[i]);
    }
    free(paths);
    return NULL; 
}


bool* isValidPath(uint8_t** path) {
    static bool validity[2];

    // Initialize to false
    validity[0] = false;
    validity[1] = false;

    // Check if path is not NULL
    if (path != NULL) {
        // Check the first element of the shortest path dimension
        if (path[0][0] != 100) {
            validity[0] = true;
        }

        // Check the first element of the reliable path dimension
        if (path[1][0] != 100) {
            validity[1] = true;
        }
    }

    return validity;
}