#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "config.h"
#include "UWBHelper.h"

// Function declarations
void showStartupScreen();
void processUWBData();
void updateDisplay();
void checkTimeouts();

// Hardware setup
HardwareSerial uwbSerial(2);
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
UWBHelper uwb(&uwbSerial, RESET_PIN);

// Data storage
DeviceData anchors[MAX_ANCHORS];
String response = "";
unsigned long lastDisplayUpdate = 0;

void setup() {
    Serial.begin(SERIAL_BAUD);
    
    // Initialize display
    Wire.begin(I2C_SDA, I2C_SCL);
    if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.println(F("SSD1306 allocation failed"));
        for (;;);
    }
    
    showStartupScreen();
    
    // Initialize UWB
    uwbSerial.begin(115200, SERIAL_8N1, IO_RXD2, IO_TXD2);
    
    if (!uwb.begin()) {
        Serial.println("UWB initialization failed!");
        return;
    }
    
    // Configure as tag
    Serial.println("Configuring as Tag...");
    uwb.configureDevice(UWB_INDEX, false); // false = tag
    uwb.setCapacity(10, 10, 1);
    uwb.setNetwork(NETWORK_ID);
    uwb.enableReporting(true);
    uwb.saveConfiguration();
    delay(1000);
    uwb.restartDevice();
    delay(2000);
    
    // Initialize anchor data
    for (int i = 0; i < MAX_ANCHORS; i++) {
        anchors[i].deviceId = -1;
        anchors[i].active = false;
    }
    
    Serial.println("Tag ready!");
}

void loop() {
    processUWBData();
    updateDisplay();
    checkTimeouts();
    
    // Handle serial passthrough for debugging
    while (Serial.available()) {
        uwbSerial.write(Serial.read());
    }
    
    delay(10);
}

void processUWBData() {
    while (uwbSerial.available()) {
        char c = uwbSerial.read();
        if (c == '\r') continue;
        if (c == '\n') {
            if (response.length() > 0) {
                Serial.println("Received: " + response);
                if (uwb.parseRangeData(response, anchors, MAX_ANCHORS)) {
                    // Successfully parsed data
                    for (int i = 0; i < MAX_ANCHORS; i++) {
                        if (anchors[i].active && (millis() - anchors[i].lastUpdate < 1000)) {
                            Serial.printf("Anchor %d: %.2fm, RSSI: %.1fdBm\n", 
                                        anchors[i].deviceId, anchors[i].distance, anchors[i].rssi);
                        }
                    }
                }
                response = "";
            }
        } else {
            response += c;
        }
    }
}

void updateDisplay() {
    if (millis() - lastDisplayUpdate < DISPLAY_UPDATE_INTERVAL) return;
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    
    // Header
    display.setCursor(0, 0);
    display.printf("Tag %d (Net:%d)", UWB_INDEX, NETWORK_ID);
    
    display.setCursor(0, 10);
    display.println("Anchors:");
    
    int activeCount = 0;
    int yPos = 20;
    
    for (int i = 0; i < MAX_ANCHORS && yPos < 55; i++) {
        if (anchors[i].active) {
            display.setCursor(0, yPos);
            display.printf("A%d: %.2fm", anchors[i].deviceId, anchors[i].distance);
            display.setCursor(70, yPos);
            display.printf("%.0fdBm", anchors[i].rssi);
            yPos += 8;
            activeCount++;
        }
    }
    
    if (activeCount == 0) {
        display.setCursor(0, 30);
        display.println("Searching...");
        display.setCursor(0, 40);
        display.println("for anchors...");
    }
    
    // Status line
    display.setCursor(0, 56);
    display.printf("Found: %d/%d", activeCount, MAX_ANCHORS);
    
    display.display();
    lastDisplayUpdate = millis();
}

void checkTimeouts() {
    unsigned long currentTime = millis();
    for (int i = 0; i < MAX_ANCHORS; i++) {
        if (anchors[i].active && (currentTime - anchors[i].lastUpdate > DEVICE_TIMEOUT)) {
            anchors[i].active = false;
        }
    }
}


void showStartupScreen() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("MaUWB Tag");
    display.setCursor(0, 15);
    display.printf("ID: %d", UWB_INDEX);
    display.setCursor(0, 25);
    display.printf("Network: %d", NETWORK_ID);
    display.setCursor(0, 35);
    display.println("Initializing...");
    display.display();
    delay(2000);
}