#include <WiFi.h>
#include <HTTPClient.h>
#include <ArduinoJson.h> // For JSON handling

const char* ssid = "Wokwi-GUEST"; // Wokwi's simulated WiFi SSID [3]
const char* password = "";        // No password for Wokwi-GUEST [3]

// !!! IMPORTANT: REPLACE WITH YOUR ACTUAL RENDER URL !!!
const char* serverUrl = "https://aiht.onrender.com/infer"; // <-- EDIT THIS LINE

const int NUM_INPUT_FEATURES = 6;  // Should match your model's input
const int NUM_OUTPUT_FEATURES = 4; // Should match your model's output

// Simulate generating input data (6 float values)
void generateModelInput(float* inputArray) {
  for (int i = 0; i < NUM_INPUT_FEATURES; i++) {
    inputArray[i] = (float)random(0, 1000) / 100.0; // Example values 0.0 to 9.99
  }
}

void setup() {
  Serial.begin(115200);
  Serial.println("ESP32 ONNX Client Started");

  // Connect to WiFi [3]
  Serial.print("Connecting to WiFi...");
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  Serial.println("\nConnected to WiFi");
  Serial.print("IP Address: ");
  Serial.println(WiFi.localIP());
  Serial.println("-----------------------------------------");

  randomSeed(analogRead(0)); // Seed random generator
}

void loop() {
  // 1. Generate input data
  float modelInput[NUM_INPUT_FEATURES];
  generateModelInput(modelInput);

  Serial.print("Generated Input: [");
  for(int i=0; i<NUM_INPUT_FEATURES; ++i) {
    Serial.print(modelInput[i], 4); // Print with 4 decimal places
    if (i < NUM_INPUT_FEATURES - 1) Serial.print(", ");
  }
  Serial.println("]");

  // 2. Prepare JSON payload
  StaticJsonDocument<256> jsonDoc; // Adjust size if needed for more complex data
  JsonArray inputs = jsonDoc.createNestedArray("inputs");
  for (int i = 0; i < NUM_INPUT_FEATURES; ++i) {
    inputs.add(modelInput[i]);
  }
  String requestBody;
  serializeJson(jsonDoc, requestBody);

  // 3. Send data to Flask server
  Serial.println("Sending data to server: " + String(serverUrl));
  HTTPClient http;
  http.begin(serverUrl);
  http.addHeader("Content-Type", "application/json");

  int httpResponseCode = http.POST(requestBody);

  // 4. Process response
  if (httpResponseCode > 0) {
    Serial.printf("HTTP Response code: %d\n", httpResponseCode);
    String responsePayload = http.getString();
    Serial.println("Response Payload: " + responsePayload);

    // Parse the JSON response
    StaticJsonDocument<256> responseDoc; // Adjust size if needed
    DeserializationError error = deserializeJson(responseDoc, responsePayload);

    if (error) {
      Serial.print("deserializeJson() failed: ");
      Serial.println(error.c_str());
    } else {
      // Check if 'outputs' key exists and is an array of the correct size
      JsonVariant outputsVariant = responseDoc["outputs"];
      if (!outputsVariant.is<JsonArray>() || outputsVariant.as<JsonArray>().size() != NUM_OUTPUT_FEATURES) {
         Serial.println("Error: 'outputs' array missing, not an array, or wrong size in response.");
      } else {
        JsonArray outputs = outputsVariant.as<JsonArray>();
        Serial.print("Received ONNX Output: [");

        float maxVal = -INFINITY;
        int maxIndex = -1;
        const char* classLabels[NUM_OUTPUT_FEATURES] = {"very dry", "dry", "wet", "very wet"};

        for (int i = 0; i < NUM_OUTPUT_FEATURES; ++i) {
          float val = outputs[i].as<float>();
          Serial.print(val, 4);
          if (val > maxVal) {
            maxVal = val;
            maxIndex = i;
          }
          if (i < NUM_OUTPUT_FEATURES - 1) Serial.print(", ");
        }
        Serial.println("]");

        if (maxIndex >= 0) {
          Serial.print("Predicted Class: ");
          Serial.println(classLabels[maxIndex]);
        }
      }
    }
  } else {
    Serial.printf("HTTP POST failed, error: %s\n", http.errorToString(httpResponseCode).c_str());
  }

  http.end(); // Free resources

  Serial.println("-----------------------------------------");
  Serial.println("Next cycle in 10 seconds...\n");
  delay(10000); // Wait 10 seconds before next cycle
}