#include <Adafruit_SSD1306.h>
#include <Wire.h>
#include <WiFi.h>
#include <HTTPClient.h>
#include <ArduinoJson.h>
// #include "secrets.h" // WiFi Configuration (WiFi name and Password)
#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels
#define OLED_RESET     -1 // Reset pin # (or -1 if sharing Arduino reset pin)
#define SCREEN_ADDRESS 0x3C ///< See datasheet for Address; 0x3D for 128x64, 0x3C for 128x32
#define BTN_PIN 5
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

const char* ssid = "Wokwi-GUEST";
const char* password = "";

const String url = "https://api.openai.com/v1/chat/completions";
const String apiKey = "<your key here>";  // Replace with your OpenAI API key

String prompt = "Namaku yogi";

void setup() {
  pinMode(BTN_PIN, INPUT_PULLUP);

  Serial.begin(115200);

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
  }

  display.clearDisplay();
  display.setTextSize(1);           
  display.setTextColor(SSD1306_WHITE);      
  display.setCursor(0,0); 
  display.println("Connecting to WiFi...");
  display.display();

  WiFi.begin(ssid, password, 6);

  while (WiFi.status() != WL_CONNECTED)
  {
    display.print(".");
    delay(500);
  }

  Serial.println("CONNECTED to SSID: ");
  Serial.print(ssid);
  Serial.println("IP: ");
  Serial.print(WiFi.localIP());

  display.clearDisplay();
  display.setCursor(0,0); 
  display.println("Connected to ");
  display.println(ssid);
  display.display();
  
  delay(1000);

  startingDisplay();
}

void loop() {
  if (Serial.available()) {
    LoadingDisplay();

    prompt = Serial.readStringUntil('\n');
    prompt.trim();

    String answer = getGPTAnswer();
    Serial.println("Answer: " + answer);
    Serial.println("Enter a prompt:");

    answerDisplay(answer);
  }
}

void startingDisplay() {
    display.clearDisplay();
    display.setTextColor(WHITE);
    display.setTextSize(1);
    display.setFont(NULL);
    display.setCursor(10, 15);
    display.println("ChatGPT 3.5 Ready!");
    display.setCursor(7, 42);
    display.println("What's your thought?");
    display.setTextWrap(0);
    display.display();
}


void LoadingDisplay() {
    display.clearDisplay();
    display.setTextColor(WHITE);
    display.setTextSize(1);
    display.setFont(NULL);
    display.setCursor(0, 10);
    display.setTextWrap(0);
    display.setCursor(31, 10);
    display.println("Loading....");
    display.display();
}


void answerDisplay(String answer) {
    display.clearDisplay();
    display.setTextWrap(1);
    display.setTextColor(WHITE);
    display.setTextSize(1);
    display.setFont(NULL);
    display.setCursor(0, 0);
    display.println("Answer:");
    display.drawLine(0, 10, 14000, 1, 1);
    display.setCursor(0, 14);
    display.println(answer);
    display.display();
}


String getGPTAnswer() {
  HTTPClient http;
  http.begin(url);
  http.addHeader("Content-Type", "application/json");
  http.addHeader("Authorization", "Bearer " + apiKey);

  String data = "\{ \"model\": \"gpt-3.5-turbo\",\"messages\": \[{\"role\": \"user\",\"content\":\""  + prompt + "\"\}\]\}";

  Serial.println("Prompt: " + prompt);

  int httpResponseCode = http.POST(data);

  if (httpResponseCode == 200) {
    String response = http.getString();
    DynamicJsonDocument doc(2048);
    DeserializationError error = deserializeJson(doc, response);

    if (error) {
      Serial.print("deserializeJson() failed: ");
      Serial.println(error.c_str());
      return "<error>";
    }

    String answer = doc["choices"][0]["message"]["content"].as<String>();;
    http.end();
    return answer;
  } else {
    Serial.print("HTTP POST request failed, error code: ");
    Serial.println(httpResponseCode);
    http.end();
    return "<error>";
  }
}