#include <WiFi.h>
#include <WiFiClientSecure.h>
#include <ArduinoJson.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// OLED display width and height
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Define OLED I2C address
#define OLED_I2C_ADDRESS 0x3C

// Create an SSD1306 display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Wi-Fi credentials for Wokwi
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// Cohere API settings
const char* cohere_host = "api.cohere.ai";
const int https_port = 443; 
const char* api_key = "4craYYzTsSPEMQtKMniL6gvoYYg1LQMpNnDo73bB"; // Replace with your API Key

WiFiClientSecure client;

int buttonPin = 2; // Button connected to GPIO2

void setup() {
  Serial.begin(115200);
  Serial.println("Starting Cohere API with OLED Display...");

  // Configure the button pin as input with pull-up resistor
  pinMode(buttonPin, INPUT_PULLUP);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_I2C_ADDRESS)) {
    Serial.println("Failed to initialize OLED display!");
    while (true); // Halt
  }

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

  // Connect to Wi-Fi
  connectToWiFi();

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Press the button to send text");
  display.display();
}

void connectToWiFi() {
  Serial.println("Connecting to Wi-Fi...");
  WiFi.begin(ssid, password);

  while (WiFi.status() != WL_CONNECTED) {
    delay(1000);
    Serial.print(".");
  }

  Serial.println("\nConnected to Wi-Fi!");
  Serial.print("IP Address: ");
  Serial.println(WiFi.localIP());

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Wi-Fi Connected!");
  display.display();
}

void sendToCohere(String user_prompt) {
  Serial.println("Connecting to Cohere server...");

  client.setInsecure();        // Disable SSL certificate verification
  client.setTimeout(15000);    // Set timeout to 15 seconds

  if (!client.connect(cohere_host, https_port)) {
    Serial.println("Failed to connect to Cohere server!");
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("Failed to connect!");
    display.display();
    return;
  }

  // Create JSON payload
  String payload = "{\"model\": \"command-xlarge\", \"prompt\": \"" + user_prompt + "\", \"max_tokens\": 50}";

  // Send HTTP POST request
  client.println("POST /generate HTTP/1.1");
  client.println("Host: api.cohere.ai");
  client.println("Authorization: Bearer " + String(api_key));
  client.println("Content-Type: application/json");
  client.println("Content-Length: " + String(payload.length()));
  client.println();
  client.print(payload);

  // Wait for response
  Serial.println("Waiting for response...");
  while (client.connected()) {
    String line = client.readStringUntil('\n');
    if (line == "\r") break;  // End of headers
    Serial.println(line);
  }

  String response = client.readString();
  Serial.println("Response:");
  Serial.println(response);

  // Parse JSON response
  DynamicJsonDocument doc(2048);
  DeserializationError error = deserializeJson(doc, response);

  if (error) {
    Serial.print("JSON Deserialization failed: ");
    Serial.println(error.c_str());
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("JSON Error!");
    display.display();
    return;
  }

  // Extract generated text
  String generated_text = doc["text"].as<String>();
  Serial.println("Generated Text:");
  Serial.println(generated_text);

  // Display text on OLED
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Generated Text:");
  display.setCursor(0, 10);
  display.println(generated_text);
  display.display();
}

void loop() {
  // Check if the button is pressed
  if (digitalRead(buttonPin) == LOW) { // Button pressed (active LOW)
    Serial.println("Button pressed! Sending text to Cohere...");
    sendToCohere("How are you?");
    delay(1000); // Debounce delay to prevent multiple triggers
  }
}