#include <WiFi.h>
#include <HTTPClient.h>
#include <ArduinoJson.h>  // Include ArduinoJson library for parsing JSON

// Replace with your Wi-Fi credentials
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// OpenAI API details
const char* apiKey = "jAd7rG1jnULkptL28GQWLOELCnW4nN0N";
const char* endpoint = "https://api.ai21.com/studio/v1/chat/completions";

String userPrompt = "";

void setup() {
  Serial.begin(115200);
  WiFi.begin(ssid, password);

  // Wait for Wi-Fi to connect
  while (WiFi.status() != WL_CONNECTED) {
    delay(1000);
    Serial.println("Connecting to Wi-Fi...");
  }
  Serial.println("Connected to Wi-Fi!");
  Serial.println("Enter your prompt below:");
}

void loop() {
  if (Serial.available()) {
    userPrompt = Serial.readStringUntil('\n'); // Read user input from Serial Monitor
    userPrompt.trim(); // Remove extra spaces or newlines

    if (userPrompt.length() > 0) {
      sendRequest(userPrompt); // Call function to send the request
    }
  }
}

void sendRequest(String prompt) {
  if (WiFi.status() == WL_CONNECTED) {
    HTTPClient http;
    http.begin(endpoint);
    http.addHeader("Content-Type", "application/json");
    http.addHeader("Authorization", String("Bearer ") + apiKey);

    // Create JSON payload
    String payload = R"({
      "model": "jamba-1.5-large",
      "messages": [{"role": "user", "content": ")" + prompt + R"("}],
      "max_tokens": 200,
      "temperature": 0.7
    })";

    int httpResponseCode = http.POST(payload);

    if (httpResponseCode > 0) {
      String response = http.getString();
      Serial.println("Raw Response:");
      Serial.println(response);

      // Parse the JSON response
      DynamicJsonDocument doc(2048); // Allocate memory for JSON parsing
      DeserializationError error = deserializeJson(doc, response);

      if (!error) {
          String content1 = doc["choices"][0]["message"]["role"].as<String>();
        String content = doc["choices"][0]["message"]["content"].as<String>();
        Serial.println("Content Response:");
        Serial.println(content1 + " : " + content);
      } else {
        Serial.print("JSON Parsing Error: ");
        Serial.println(error.c_str());
      }
    } else {
      Serial.print("Error on sending POST: ");
      Serial.println(httpResponseCode);
    }
    http.end();
  } else {
    Serial.println("Wi-Fi not connected!");
  }
}