#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 tags[UWB_TAG_COUNT];
DistanceFilter tagFilters[UWB_TAG_COUNT];
String response = "";
unsigned long lastDisplayUpdate = 0;

void setup() {
    Serial.begin(SERIAL_BAUD);
    Serial.println("Starting UWB Anchor...");
    
    // Initialize I2C and 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);
    delay(1000);
    
    if (!uwb.begin()) {
        Serial.println("UWB initialization failed!");
        return;
    }
    
    // Configure as anchor
    Serial.println("Configuring as Anchor...");
    uwb.configureDevice(UWB_INDEX, true); // true = anchor
    uwb.setCapacity(UWB_TAG_COUNT, 10, 1);
    uwb.setNetwork(NETWORK_ID);
    uwb.enableReporting(true);
    uwb.saveConfiguration();
    delay(1000);
    uwb.restartDevice();
    delay(2000);
    
    // Initialize tag data
    for (int i = 0; i < UWB_TAG_COUNT; i++) {
        tags[i].deviceId = -1;
        tags[i].active = false;
        tags[i].distance = 0.0;
        tags[i].rssi = 0.0;
        tags[i].lastUpdate = 0;
    }
    
    Serial.println("Anchor ready! Waiting for tags...");
}

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, tags, UWB_TAG_COUNT)) {
                    // Successfully parsed data
                    for (int i = 0; i < UWB_TAG_COUNT; i++) {
                        if (tags[i].active && (millis() - tags[i].lastUpdate < 1000)) {
                            Serial.printf("Tag %d: %.2fm, RSSI: %.1fdBm\n", 
                                        tags[i].deviceId, tags[i].distance, tags[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("Anchor %d (Net:%d)", UWB_INDEX, NETWORK_ID);
    
    display.setCursor(0, 10);
    display.println("Active Tags:");
    
    // Show active tags
    int activeCount = 0;
    int yPos = 20;
    
    for (int i = 0; i < UWB_TAG_COUNT && yPos < 55; i++) {
        if (tags[i].active) {
            display.setCursor(0, yPos);
            display.printf("T%d: %.2fm", tags[i].deviceId, tags[i].distance);
            display.setCursor(70, yPos);
            display.printf("%.0fdBm", tags[i].rssi);
            yPos += 8;
            activeCount++;
        }
    }
    
    if (activeCount == 0) {
        display.setCursor(0, 30);
        display.println("No active tags");
        display.setCursor(0, 40);
        display.println("Waiting...");
    }
    
    // Status line
    display.setCursor(0, 56);
    display.printf("Active: %d/%d", activeCount, UWB_TAG_COUNT);
    
    display.display();
    lastDisplayUpdate = millis();
}

void checkTimeouts() {
    unsigned long currentTime = millis();
    for (int i = 0; i < UWB_TAG_COUNT; i++) {
        if (tags[i].active && (currentTime - tags[i].lastUpdate > DEVICE_TIMEOUT)) {
            tags[i].active = false;
            Serial.printf("Tag %d timeout\n", tags[i].deviceId);
        }
    }
}

void showStartupScreen() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("MaUWB Anchor");
    
    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("6.8Mbps Mode");
    
    display.setCursor(0, 45);
    display.println("Initializing...");
    
    display.display();
    delay(2000);
}