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

// Function declarations
void showStartupScreen();
void processUWBData();
void updateDisplay();
void checkTimeouts();
void processAnchorBehavior();
void processTagBehavior();
void handleSerialPassthrough();
void updateBatteryLevel();
void drawBatteryIcon(int x, int y, int percent);

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

// Get device configuration
DeviceConfig deviceConfig = getDeviceConfig();

// Data storage based on device type
#if DEVICE_TYPE == DEVICE_ANCHOR
DeviceData tags[UWB_TAG_COUNT];
DistanceFilter tagFilters[UWB_TAG_COUNT];
#elif DEVICE_TYPE == DEVICE_TAG  
DeviceData anchors[MAX_ANCHORS];
#endif

String response = "";
unsigned long lastDisplayUpdate = 0;
unsigned long lastPositioningUpdate = 0;
unsigned long lastBatteryUpdate = 0;
float batteryVoltage = 0.0;
int batteryPercent = 0;

void setup() {
    Serial.begin(SERIAL_BAUD);
    Serial.printf("Starting UWB %s (ID: %d)...\n", deviceConfig.deviceName, deviceConfig.deviceId);
    
    // Initialize display
    Wire.end(); // Ensure I2C bus is properly reset
    delay(100);
    Wire.begin(I2C_SDA, I2C_SCL);
    if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.println(F("SSD1306 allocation failed"));
        for (;;);
    }
    
    showStartupScreen();
    
    // Initialize UWB serial
    uwbSerial.begin(115200, SERIAL_8N1, IO_RXD2, IO_TXD2);
    
    // Hardware reset UWB module (critical for reliable operation)
    pinMode(RESET_PIN, OUTPUT);
    digitalWrite(RESET_PIN, HIGH);
    delay(1000);
    
    Serial.println("Configuring UWB module...");
    
    // Complete UWB setup sequence following AT command protocol
    uwb.sendCommand("AT?", 2000);
    delay(500);
    
    uwb.sendCommand("AT+RESTORE", 5000);
    delay(2000);  // Wait for factory reset to complete
    
    // Configure device role and parameters
    String configCmd = "AT+SETCFG=" + String(deviceConfig.deviceId) + ",";
    configCmd += deviceConfig.isAnchor ? "1" : "0";  // 1=Anchor, 0=Tag
    configCmd += ",1,1";  // 6.8Mbps mode, range filtering enabled
    uwb.sendCommand(configCmd, 2000);
    delay(500);
    
    // Set device capacity
    String capacityCmd = "AT+SETCAP=" + String(deviceConfig.capacityTags) + ",";
    capacityCmd += String(deviceConfig.capacityTimeSlot) + ",";
    capacityCmd += String(deviceConfig.capacityExtMode);
    uwb.sendCommand(capacityCmd, 2000);
    delay(500);
    
    // Enable automatic reporting and set network
    uwb.sendCommand("AT+SETRPT=1", 2000);
    delay(500);
    uwb.sendCommand("AT+SETPAN=" + String(NETWORK_ID), 2000);
    delay(500);
    
    // Save configuration and restart
    uwb.sendCommand("AT+SAVE", 2000);
    delay(1000);  // Wait for save to complete
    uwb.sendCommand("AT+RESTART", 2000);
    delay(3000);  // Wait longer for restart to complete
    
    // Verify configuration
    Serial.println("Verifying configuration...");
    String configResponse = uwb.sendCommand("AT+GETCFG?", 2000);
    Serial.printf("Configuration verification: %s\n", configResponse.c_str());
    
    // Additional verification commands
    uwb.sendCommand("AT+GETPAN?", 2000);
    uwb.sendCommand("AT+GETRPT?", 2000);
    
    // Initialize device data
    #if DEVICE_TYPE == DEVICE_ANCHOR
        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;
        }
    #elif DEVICE_TYPE == DEVICE_TAG
        for (int i = 0; i < MAX_ANCHORS; i++) {
            anchors[i].deviceId = -1;
            anchors[i].active = false;
            anchors[i].distance = 0.0;
            anchors[i].rssi = 0.0;
            anchors[i].lastUpdate = 0;
        }
    #endif
    
    // Initial battery reading
    updateBatteryLevel();
    
    Serial.printf("%s ready! Network ID: %d\n", deviceConfig.deviceName, NETWORK_ID);
    
    #if DEVICE_TYPE == DEVICE_ANCHOR
        Serial.println("Anchor behaviors enabled:");
        #ifdef ENABLE_TAG_TRACKING
            Serial.println("- Tag tracking");
        #endif
        #ifdef ENABLE_INTER_ANCHOR_COMM
            Serial.println("- Inter-anchor communication");
        #endif
        #ifdef ENABLE_POSITION_CALCULATION
            Serial.println("- Position calculation");
        #endif
    #elif DEVICE_TYPE == DEVICE_TAG
        Serial.println("Tag behaviors enabled:");
        #ifdef ENABLE_ANCHOR_TRACKING
            Serial.println("- Anchor tracking");
        #endif
        #ifdef ENABLE_USB_DATA_STREAMING
            Serial.println("- USB data streaming");
        #endif
        #ifdef ENABLE_COORDINATE_COLLECTION
            Serial.println("- Coordinate collection");
        #endif
    #endif
}

void loop() {
    processUWBData();
    updateDisplay();
    checkTimeouts();
    handleSerialPassthrough();
    
    // Update battery level periodically
    if (millis() - lastBatteryUpdate > BATTERY_UPDATE_INTERVAL) {
        updateBatteryLevel();
        lastBatteryUpdate = millis();
    }
    
    // Execute device-specific behavior
    if (DEVICE_IS_ANCHOR) {
        processAnchorBehavior();
    } else {
        processTagBehavior();
    }
    
    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 DEVICE_TYPE == DEVICE_ANCHOR
                    if (uwb.parseRangeData(response, tags, UWB_TAG_COUNT, true)) {
                        // Successfully parsed data - log active tags
                        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);
                            }
                        }
                    }
                #elif DEVICE_TYPE == DEVICE_TAG
                    if (uwb.parseRangeData(response, anchors, MAX_ANCHORS, false)) {
                        // Successfully parsed data - log active anchors
                        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);
                            }
                        }
                    }
                #endif
                response = "";
            }
        } else {
            response += c;
        }
    }
}

void updateDisplay() {
    if (millis() - lastDisplayUpdate < DISPLAY_UPDATE_INTERVAL) return;
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    
    // Header - Device info with battery
    display.setCursor(0, 0);
    display.printf("%s %d", deviceConfig.deviceName, deviceConfig.deviceId);
    
    // Battery indicator in top right
    drawBatteryIcon(80, 0, batteryPercent);
    
    // Status line
    display.setCursor(0, 10);
    display.println(deviceConfig.statusDisplayPrefix);
    
    // Show active peer devices
    int activeCount = 0;
    int yPos = 20;
    
    #if DEVICE_TYPE == DEVICE_ANCHOR
        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++;
            }
        }
    #elif DEVICE_TYPE == DEVICE_TAG
        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++;
            }
        }
    #endif
    
    if (activeCount == 0) {
        display.setCursor(0, 30);
        display.println(deviceConfig.waitingMessage);
        display.setCursor(0, 40);
        display.println(deviceConfig.searchingMessage);
    }
    
    // Status line with device-specific information
    display.setCursor(0, 56);
    if (DEVICE_IS_ANCHOR) {
        display.printf("Active: %d/%d", activeCount, UWB_TAG_COUNT);
    } else {
        display.printf("Found: %d/%d", activeCount, MAX_ANCHORS);
    }
    
    display.display();
    lastDisplayUpdate = millis();
}

void checkTimeouts() {
    unsigned long currentTime = millis();
    
    #if DEVICE_TYPE == DEVICE_ANCHOR
        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);
            }
        }
    #elif DEVICE_TYPE == DEVICE_TAG
        for (int i = 0; i < MAX_ANCHORS; i++) {
            if (anchors[i].active && (currentTime - anchors[i].lastUpdate > DEVICE_TIMEOUT)) {
                anchors[i].active = false;
                Serial.printf("Anchor %d timeout\n", anchors[i].deviceId);
            }
        }
    #endif
}

void processAnchorBehavior() {
    #if DEVICE_TYPE == DEVICE_ANCHOR
        unsigned long currentTime = millis();
        
        #ifdef ENABLE_INTER_ANCHOR_COMM
        // TODO: Implement inter-anchor discovery and communication
        static unsigned long lastDiscovery = 0;
        if (currentTime - lastDiscovery > ANCHOR_DISCOVERY_INTERVAL) {
            // Broadcast discovery signal to other anchors
            DEBUG_PRINTLN("Broadcasting anchor discovery...");
            lastDiscovery = currentTime;
        }
        #endif
        
        #ifdef ENABLE_POSITION_CALCULATION
        // TODO: Implement distributed position calculation
        if (currentTime - lastPositioningUpdate > POSITION_UPDATE_INTERVAL) {
            // Calculate and update anchor position
            DEBUG_PRINTLN("Updating anchor position...");
            lastPositioningUpdate = currentTime;
        }
        #endif
        
        #ifdef ENABLE_COORDINATE_STORAGE
        // TODO: Store calculated position in EEPROM
        #endif
    #endif
}

void processTagBehavior() {
    #if DEVICE_TYPE == DEVICE_TAG
        unsigned long currentTime = millis();
        
        #ifdef ENABLE_USB_DATA_STREAMING
        // Stream positioning data via USB
        static unsigned long lastStream = 0;
        if (currentTime - lastStream > USB_STREAM_INTERVAL) {
            // TODO: Stream formatted positioning data
            lastStream = currentTime;
        }
        #endif
        
        #ifdef ENABLE_COORDINATE_COLLECTION
        // Request anchor positions
        static unsigned long lastPositionRequest = 0;
        if (currentTime - lastPositionRequest > POSITION_REQUEST_INTERVAL) {
            // TODO: Request positions from connected anchors
            DEBUG_PRINTLN("Requesting anchor positions...");
            lastPositionRequest = currentTime;
        }
        #endif
        
        #ifdef ENABLE_REAL_TIME_POSITIONING
        // Calculate tag position
        if (currentTime - lastPositioningUpdate > POSITION_UPDATE_RATE_MS) {
            // TODO: Calculate real-time tag position
            lastPositioningUpdate = currentTime;
        }
        #endif
    #endif
}

void handleSerialPassthrough() {
    // Handle serial passthrough for debugging
    while (Serial.available()) {
        uwbSerial.write(Serial.read());
    }
}

void showStartupScreen() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.printf("MaUWB %s", deviceConfig.deviceName);
    
    display.setCursor(0, 15);
    display.printf("ID: %d", deviceConfig.deviceId);
    
    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);
}

void updateBatteryLevel() {
    batteryVoltage = 3.3 * analogRead(BAT_PIN) / 4096.0 * 2;
    batteryPercent = constrain(map((int)(batteryVoltage * 100), 300, 420, 0, 100), 0, 100);
}

void drawBatteryIcon(int x, int y, int percent) {
    display.drawRect(x, y, 18, 10, SSD1306_WHITE);
    display.fillRect(x + 18, y + 3, 2, 4, SSD1306_WHITE);
    
    int fillWidth = (14 * percent) / 100;
    if (fillWidth > 0) {
        display.fillRect(x + 2, y + 2, fillWidth, 6, SSD1306_WHITE);
    }
    
    display.setCursor(x + 22, y + 1);
    display.setTextSize(1);
    display.printf("%d%%", percent);
}

// Cleanup - no dynamic allocation needed