Desk Voice Assistant

/**
 * 🤖 AI Desk Assistant - ESP32-S3 DevKitC-1 + GC9A01 Round Display
 *
 * Flow:
 *   1. Press BOOT button → record 5 s from INMP441
 *   2. Whisper STT → transcript
 *   3. GPT-4o-mini → reply text (shown on display)
 *   4. OpenAI TTS → MP3 → MAX98357A speaker
 *
 * ── PIN MAP ──────────────────────────────────────────
 *  INMP441  SCK  → GPIO1    WS → GPIO2    SD → GPIO3
 *  MAX98357 BCLK → GPIO5   LRC → GPIO6   DIN → GPIO7
 *  GC9A01   SCK  → GPIO12  MOSI→ GPIO11  CS  → GPIO10
 *           DC   → GPIO14  RST → GPIO21  BL  → 3V3
 * ─────────────────────────────────────────────────────
 *
 * IMPORTANT: Set your Wi-Fi credentials and OpenAI API key below.
 */

#include <Arduino.h>
#include <WiFi.h>
#include <HTTPClient.h>
#include <WiFiClientSecure.h>
#include <ArduinoJson.h>
#include <driver/i2s.h>
#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_GC9A01A.h>

// ── USER CONFIG ───────────────────────────────────────

struct Message { String role; String content; };

enum State { S_IDLE, S_RECORDING, S_THINKING, S_SPEAKING };

const char* WIFI_SSID     = "YOUR_WIFI_SSID";
const char* WIFI_PASSWORD = "YOUR_WIFI_PASSWORD";
const char* OPENAI_KEY    = "sk-YOUR_OPENAI_API_KEY";
// ─────────────────────────────────────────────────────

// ── Pin definitions ───────────────────────────────────
#define MIC_SCK     1
#define MIC_WS      2
#define MIC_SD      3
#define AMP_BCLK    5
#define AMP_LRC     6
#define AMP_DIN     7

#define TFT_SCK    12
#define TFT_MOSI   11
#define TFT_CS     10
#define TFT_DC     14
#define TFT_RST    21

#define BUTTON_PIN 0   // Onboard BOOT button (active LOW)

// ── I2S ports ─────────────────────────────────────────
#define I2S_MIC_PORT  I2S_NUM_0
#define I2S_AMP_PORT  I2S_NUM_1

// ── Recording settings ────────────────────────────────
#define SAMPLE_RATE        16000
#define SAMPLE_BITS        16
#define RECORD_SECONDS     5
#define PCM_BUFFER_SIZE    (SAMPLE_RATE * (SAMPLE_BITS / 8) * RECORD_SECONDS)

// ── Playback settings ─────────────────────────────────
#define PLAYBACK_RATE     24000
#define MP3_BUFFER_SIZE   (128 * 1024)

// ── Conversation history ──────────────────────────────

static const int MAX_HISTORY = 8;
Message history[MAX_HISTORY];
int historyCount = 0;

// ── Audio buffers ─────────────────────────────────────
int16_t* pcmBuffer = nullptr;
uint8_t* mp3Buffer = nullptr;

// ── Display ───────────────────────────────────────────
Adafruit_GC9A01A tft(TFT_CS, TFT_DC, TFT_MOSI, TFT_SCK, TFT_RST);

// GC9A01 is 240×240; centre is (120,120), radius = 120
#define DISPLAY_W   240
#define DISPLAY_H   240
#define CX          120
#define CY          120

// 16-bit RGB565 colour helpers
#define COL_BG        0x0841   // Very dark navy
#define COL_RING      0x04FF   // Cyan ring
#define COL_RING_REC  0xF800   // Red ring when recording
#define COL_RING_THINK 0xFFE0  // Yellow ring when thinking
#define COL_RING_SPEAK 0x07E0  // Green ring when speaking
#define COL_TEXT      0xFFFF   // White
#define COL_SUBTEXT   0x7BEF   // Light grey
#define COL_FACE_BG   0x1082   // Dark blue-grey face

// ── State machine ─────────────────────────────────────

State curState = S_IDLE;

// ── Forward declarations ──────────────────────────────
void drawFace(State s, const String& label, const String& sub = "");
void animateRing(uint16_t colour, int frames);
void setupMicI2S();
void teardownMicI2S();
void setupAmpI2S();
void teardownAmpI2S();
bool recordAudio();
void writeWavHeader(uint8_t* buf, uint32_t pcmBytes, uint32_t sr, uint16_t ch, uint16_t bps);
String whisperTranscribe(int16_t* samples, size_t sampleCount);
String chatCompletion(const String& userMessage);
bool ttsSpeak(const String& text);
void playMP3(const uint8_t* data, size_t length);
void pushHistory(const String& role, const String& content);
void wrapText(const String& text, int maxChars, String& line1, String& line2);

// ─────────────────────────────────────────────────────
// ── Display helpers ───────────────────────────────────
// ─────────────────────────────────────────────────────

// Draw the full face for a given state
void drawFace(State s, const String& label, const String& sub) {
    // Background circle
    tft.fillCircle(CX, CY, 119, COL_FACE_BG);

    // Outer ring colour by state
    uint16_t ringCol = COL_RING;
    if (s == S_RECORDING) ringCol = COL_RING_REC;
    else if (s == S_THINKING) ringCol = COL_RING_THINK;
    else if (s == S_SPEAKING) ringCol = COL_RING_SPEAK;

    // Draw 4-pixel-wide ring
    for (int r = 119; r > 115; r--) {
        tft.drawCircle(CX, CY, r, ringCol);
    }

    // Robot eyes (two rounded rects)
    int eyeY = 85;
    int eyeW = 28, eyeH = 18, eyeR = 5;
    tft.fillRoundRect(CX - 40 - eyeW/2, eyeY, eyeW, eyeH, eyeR, COL_RING);
    tft.fillRoundRect(CX + 40 - eyeW/2, eyeY, eyeW, eyeH, eyeR, COL_RING);

    // Mouth bar (varies by state)
    int mouthY = 140;
    if (s == S_IDLE) {
        // Smile - simple arc approximated with lines
        for (int x = -28; x <= 28; x++) {
            int y = (x * x) / 60;
            tft.drawPixel(CX + x, mouthY + y, COL_RING);
            tft.drawPixel(CX + x, mouthY + y + 1, COL_RING);
        }
    } else if (s == S_RECORDING) {
        // Open O mouth
        tft.drawCircle(CX, mouthY + 8, 14, COL_RING_REC);
        tft.drawCircle(CX, mouthY + 8, 13, COL_RING_REC);
    } else if (s == S_THINKING) {
        // Flat thinking line
        tft.fillRoundRect(CX - 25, mouthY + 4, 50, 6, 3, COL_RING_THINK);
    } else if (s == S_SPEAKING) {
        // Sound wave bars
        int barW = 6, barGap = 4;
        int bars[] = {10, 20, 14, 22, 10};
        int totalW = 5 * barW + 4 * barGap;
        int startX = CX - totalW / 2;
        for (int b = 0; b < 5; b++) {
            int bh = bars[b];
            int bx = startX + b * (barW + barGap);
            tft.fillRoundRect(bx, mouthY + 10 - bh / 2, barW, bh, 2, COL_RING_SPEAK);
        }
    }

    // Label text (state name)
    tft.setTextColor(COL_TEXT);
    tft.setTextSize(2);
    int labelX = CX - (label.length() * 6);
    tft.setCursor(labelX, 175);
    tft.print(label);

    // Sub text (transcript / reply preview)
    if (sub.length() > 0) {
        String l1, l2;
        wrapText(sub, 18, l1, l2);
        tft.setTextColor(COL_SUBTEXT);
        tft.setTextSize(1);
        tft.setCursor(CX - l1.length() * 3, 198);
        tft.print(l1);
        if (l2.length() > 0) {
            tft.setCursor(CX - l2.length() * 3, 210);
            tft.print(l2);
        }
    }
}

// Wrap text to two lines of maxChars each
void wrapText(const String& text, int maxChars, String& line1, String& line2) {
    if ((int)text.length() <= maxChars) {
        line1 = text; line2 = "";
        return;
    }
    int split = text.lastIndexOf(' ', maxChars);
    if (split < 0) split = maxChars;
    line1 = text.substring(0, split);
    line2 = text.substring(split + 1);
    if ((int)line2.length() > maxChars) line2 = line2.substring(0, maxChars - 1) + "...";
}

// ─────────────────────────────────────────────────────
// ── Setup ─────────────────────────────────────────────
// ─────────────────────────────────────────────────────
void setup() {
    Serial.begin(115200);
    delay(800);

    pinMode(BUTTON_PIN, INPUT_PULLUP);

    // ── Init display ──
    tft.begin();
    tft.setRotation(0);
    tft.fillScreen(0x0000);
    drawFace(S_IDLE, "Booting", "Connecting...");

    // ── Allocate audio buffers ──
    pcmBuffer = (int16_t*)heap_caps_malloc(PCM_BUFFER_SIZE, MALLOC_CAP_DEFAULT);
    mp3Buffer = (uint8_t*)heap_caps_malloc(MP3_BUFFER_SIZE, MALLOC_CAP_DEFAULT);
    if (!pcmBuffer || !mp3Buffer) {
        tft.fillCircle(CX, CY, 119, 0xF800);
        tft.setCursor(60, CY - 8);
        tft.setTextColor(COL_TEXT); tft.setTextSize(2);
        tft.print("MEM ERR");
        while (true) delay(1000);
    }

    // ── Connect Wi-Fi ──
    Serial.printf("📶 Connecting to %s...\n", WIFI_SSID);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 40) {
        delay(500); Serial.print(".");
        attempts++;
    }
    if (WiFi.status() != WL_CONNECTED) {
        drawFace(S_IDLE, "No WiFi", "Check credentials");
        while (true) delay(1000);
    }
    Serial.printf("\n✅ IP: %s\n", WiFi.localIP().toString().c_str());

    // ── System prompt ──
    pushHistory("system",
        "You are a friendly, witty AI desk assistant. "
        "Keep replies concise - two or three sentences max. "
        "You are running on an ESP32-S3 DevKitC-1 microcontroller on someone's desk. "
        "Be helpful, a little playful, and never condescending.");

    drawFace(S_IDLE, "Ready!", "Hold BOOT to speak");
    Serial.println("🤖 Ready - press BOOT to speak!");
}

// ─────────────────────────────────────────────────────
// ── Main loop ─────────────────────────────────────────
// ─────────────────────────────────────────────────────
void loop() {
    if (digitalRead(BUTTON_PIN) == LOW) {
        delay(50);
        if (digitalRead(BUTTON_PIN) == LOW) {

            // ── RECORDING ──
            curState = S_RECORDING;
            drawFace(S_RECORDING, "Listening", "");
            Serial.println("\n🎙️  Recording...");

            setupMicI2S();
            bool gotAudio = recordAudio();
            teardownMicI2S();

            if (!gotAudio) {
                drawFace(S_IDLE, "Ready!", "Mic error-retry");
                return;
            }

            // ── THINKING: STT ──
            curState = S_THINKING;
            drawFace(S_THINKING, "Thinking", "Transcribing...");
            Serial.println("🧠 Whisper STT...");

            size_t sampleCount = PCM_BUFFER_SIZE / sizeof(int16_t);
            String transcript = whisperTranscribe(pcmBuffer, sampleCount);
            if (transcript.isEmpty()) {
                drawFace(S_IDLE, "Ready!", "Couldn't hear you");
                return;
            }
            Serial.printf("💬 You: %s\n", transcript.c_str());
            drawFace(S_THINKING, "Thinking", transcript);

            // ── THINKING: Chat ──
            Serial.println("💭 GPT...");
            String reply = chatCompletion(transcript);
            if (reply.isEmpty()) {
                drawFace(S_IDLE, "Ready!", "No reply");
                return;
            }
            Serial.printf("🤖 AI: %s\n", reply.c_str());

            // ── SPEAKING ──
            curState = S_SPEAKING;
            drawFace(S_SPEAKING, "Speaking", reply);
            Serial.println("🔊 Speaking...");

            setupAmpI2S();
            ttsSpeak(reply);
            teardownAmpI2S();

            // ── Back to idle ──
            curState = S_IDLE;
            drawFace(S_IDLE, "Ready!", "Hold BOOT to speak");

            while (digitalRead(BUTTON_PIN) == LOW) delay(10);
        }
    }
    delay(20);
}

// ─────────────────────────────────────────────────────
// ── I2S ───────────────────────────────────────────────
// ─────────────────────────────────────────────────────
void setupMicI2S() {
    i2s_config_t cfg = {
        .mode                 = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_RX),
        .sample_rate          = SAMPLE_RATE,
        .bits_per_sample      = I2S_BITS_PER_SAMPLE_32BIT,
        .channel_format       = I2S_CHANNEL_FMT_ONLY_LEFT,
        .communication_format = I2S_COMM_FORMAT_STAND_I2S,
        .intr_alloc_flags     = ESP_INTR_FLAG_LEVEL1,
        .dma_buf_count        = 8,
        .dma_buf_len          = 1024,
        .use_apll             = false,
        .tx_desc_auto_clear   = false,
        .fixed_mclk           = 0
    };
    i2s_pin_config_t pins = {
        .bck_io_num   = MIC_SCK,
        .ws_io_num    = MIC_WS,
        .data_out_num = I2S_PIN_NO_CHANGE,
        .data_in_num  = MIC_SD
    };
    i2s_driver_install(I2S_MIC_PORT, &cfg, 0, NULL);
    i2s_set_pin(I2S_MIC_PORT, &pins);
    i2s_zero_dma_buffer(I2S_MIC_PORT);
}

void teardownMicI2S() { i2s_driver_uninstall(I2S_MIC_PORT); }

void setupAmpI2S() {
    i2s_config_t cfg = {
        .mode                 = (i2s_mode_t)(I2S_MODE_MASTER | I2S_MODE_TX),
        .sample_rate          = PLAYBACK_RATE,
        .bits_per_sample      = I2S_BITS_PER_SAMPLE_16BIT,
        .channel_format       = I2S_CHANNEL_FMT_ONLY_LEFT,
        .communication_format = I2S_COMM_FORMAT_STAND_I2S,
        .intr_alloc_flags     = ESP_INTR_FLAG_LEVEL1,
        .dma_buf_count        = 8,
        .dma_buf_len          = 1024,
        .use_apll             = false,
        .tx_desc_auto_clear   = true,
        .fixed_mclk           = 0
    };
    i2s_pin_config_t pins = {
        .bck_io_num   = AMP_BCLK,
        .ws_io_num    = AMP_LRC,
        .data_out_num = AMP_DIN,
        .data_in_num  = I2S_PIN_NO_CHANGE
    };
    i2s_driver_install(I2S_AMP_PORT, &cfg, 0, NULL);
    i2s_set_pin(I2S_AMP_PORT, &pins);
    i2s_zero_dma_buffer(I2S_AMP_PORT);
}

void teardownAmpI2S() { i2s_driver_uninstall(I2S_AMP_PORT); }

// ─────────────────────────────────────────────────────
// ── Record audio ──────────────────────────────────────
// ─────────────────────────────────────────────────────
bool recordAudio() {
    size_t bytesRead = 0;
    size_t totalSamples = 0;
    const size_t maxSamples = PCM_BUFFER_SIZE / sizeof(int16_t);

    // Warm-up
    int32_t dummy[512];
    for (int w = 0; w < 5; w++) {
        i2s_read(I2S_MIC_PORT, dummy, sizeof(dummy), &bytesRead, portMAX_DELAY);
    }

    int32_t rawBuf[512];
    while (totalSamples < maxSamples) {
        i2s_read(I2S_MIC_PORT, rawBuf, sizeof(rawBuf), &bytesRead, portMAX_DELAY);
        size_t frames = bytesRead / sizeof(int32_t);
        for (size_t i = 0; i < frames && totalSamples < maxSamples; i++) {
            pcmBuffer[totalSamples++] = (int16_t)(rawBuf[i] >> 14);
        }
    }
    return totalSamples > 0;
}

// ─────────────────────────────────────────────────────
// ── WAV header ────────────────────────────────────────
// ─────────────────────────────────────────────────────
void writeWavHeader(uint8_t* buf, uint32_t pcmBytes, uint32_t sr, uint16_t ch, uint16_t bps) {
    uint32_t byteRate   = sr * ch * bps / 8;
    uint16_t blockAlign = ch * bps / 8;
    uint32_t riffSize   = 36 + pcmBytes;
    uint32_t fmtSize    = 16;
    uint16_t audioFmt   = 1;
    memcpy(buf,    "RIFF", 4); memcpy(buf+4,  &riffSize,   4);
    memcpy(buf+8,  "WAVE", 4); memcpy(buf+12, "fmt ",       4);
    memcpy(buf+16, &fmtSize,   4); memcpy(buf+20, &audioFmt,   2);
    memcpy(buf+22, &ch,        2); memcpy(buf+24, &sr,         4);
    memcpy(buf+28, &byteRate,  4); memcpy(buf+32, &blockAlign, 2);
    memcpy(buf+34, &bps,       2); memcpy(buf+36, "data",      4);
    memcpy(buf+40, &pcmBytes,  4);
}

// ─────────────────────────────────────────────────────
// ── Whisper STT ───────────────────────────────────────
// ─────────────────────────────────────────────────────
String whisperTranscribe(int16_t* samples, size_t sampleCount) {
    const int WAV_HEADER_SIZE = 44;
    uint32_t pcmBytes = sampleCount * sizeof(int16_t);
    uint8_t* wavBuf = mp3Buffer;
    if (pcmBytes + WAV_HEADER_SIZE > MP3_BUFFER_SIZE) { return ""; }

    writeWavHeader(wavBuf, pcmBytes, SAMPLE_RATE, 1, 16);
    memcpy(wavBuf + WAV_HEADER_SIZE, samples, pcmBytes);
    size_t wavSize = WAV_HEADER_SIZE + pcmBytes;

    WiFiClientSecure client; client.setInsecure();
    HTTPClient http;
    http.begin(client, "https://api.openai.com/v1/audio/transcriptions");
    http.addHeader("Authorization", String("Bearer ") + OPENAI_KEY);

    String boundary = "ESP32AI" + String(millis());
    http.addHeader("Content-Type", "multipart/form-data; boundary=" + boundary);

    String partHeader =
        "--" + boundary + "\r\n"
        "Content-Disposition: form-data; name=\"file\"; filename=\"audio.wav\"\r\n"
        "Content-Type: audio/wav\r\n\r\n";
    String modelPart =
        "\r\n--" + boundary + "\r\n"
        "Content-Disposition: form-data; name=\"model\"\r\n\r\n"
        "whisper-1\r\n"
        "--" + boundary + "--\r\n";

    size_t bodySize = partHeader.length() + wavSize + modelPart.length();
    uint8_t* bodyBuf = (uint8_t*)malloc(bodySize);
    if (!bodyBuf) { return ""; }

    memcpy(bodyBuf, partHeader.c_str(), partHeader.length());
    memcpy(bodyBuf + partHeader.length(), wavBuf, wavSize);
    memcpy(bodyBuf + partHeader.length() + wavSize, modelPart.c_str(), modelPart.length());

    int code = http.POST(bodyBuf, bodySize);
    free(bodyBuf);

    String transcript = "";
    if (code == 200) {
        JsonDocument doc;
        if (deserializeJson(doc, http.getString()) == DeserializationError::Ok) {
            transcript = doc["text"].as<String>();
            transcript.trim();
        }
    } else {
        Serial.printf("⚠️  Whisper %d\n", code);
    }
    http.end();
    return transcript;
}

// ─────────────────────────────────────────────────────
// ── Chat completion ───────────────────────────────────
// ─────────────────────────────────────────────────────
String chatCompletion(const String& userMessage) {
    pushHistory("user", userMessage);

    JsonDocument req;
    req["model"] = "gpt-4o-mini";
    JsonArray msgs = req["messages"].to<JsonArray>();
    for (int i = 0; i < historyCount; i++) {
        JsonObject m = msgs.add<JsonObject>();
        m["role"]    = history[i].role;
        m["content"] = history[i].content;
    }
    req["max_tokens"] = 150;
    String body; serializeJson(req, body);

    WiFiClientSecure client; client.setInsecure();
    HTTPClient http;
    http.begin(client, "https://api.openai.com/v1/chat/completions");
    http.addHeader("Authorization", String("Bearer ") + OPENAI_KEY);
    http.addHeader("Content-Type", "application/json");
    http.setTimeout(15000);

    int code = http.POST(body);
    String reply = "";
    if (code == 200) {
        JsonDocument doc;
        if (deserializeJson(doc, http.getString()) == DeserializationError::Ok) {
            reply = doc["choices"][0]["message"]["content"].as<String>();
            reply.trim();
            pushHistory("assistant", reply);
        }
    } else {
        Serial.printf("⚠️  Chat %d\n", code);
    }
    http.end();
    return reply;
}

// ─────────────────────────────────────────────────────
// ── TTS + playback ────────────────────────────────────
// ─────────────────────────────────────────────────────
bool ttsSpeak(const String& text) {
    JsonDocument req;
    req["model"]           = "tts-1";
    req["input"]           = text;
    req["voice"]           = "nova";
    req["response_format"] = "mp3";
    String body; serializeJson(req, body);

    WiFiClientSecure client; client.setInsecure();
    HTTPClient http;
    http.begin(client, "https://api.openai.com/v1/audio/speech");
    http.addHeader("Authorization", String("Bearer ") + OPENAI_KEY);
    http.addHeader("Content-Type", "application/json");
    http.setTimeout(15000);

    int code = http.POST(body);
    if (code == 200) {
        size_t mp3Len = http.getSize();
        if (mp3Len == 0 || mp3Len > MP3_BUFFER_SIZE) {
            http.end(); return false;
        }
        WiFiClient* stream = http.getStreamPtr();
        size_t received = 0;
        while (received < mp3Len && http.connected()) {
            int avail = stream->available();
            if (avail > 0) {
                int toRead = min((int)(MP3_BUFFER_SIZE - received), avail);
                received += stream->readBytes(mp3Buffer + received, toRead);
            }
            delay(1);
        }
        http.end();
        if (received > 0) { playMP3(mp3Buffer, received); return true; }
    } else {
        Serial.printf("⚠️  TTS %d\n", code);
        http.end();
    }
    return false;
}

// ─────────────────────────────────────────────────────
// ── MP3 decode + play ─────────────────────────────────
// ─────────────────────────────────────────────────────
#include "AudioTools.h"
#include "AudioCodecs/CodecMP3Helix.h"

void playMP3(const uint8_t* data, size_t length) {
    MemoryStream mp3Stream((uint8_t*)data, length);
    MP3DecoderHelix decoder;
    I2SStream i2sOut;
    I2SConfig i2sCfg = i2sOut.defaultConfig(TX_MODE);
    i2sCfg.port_no        = 1;
    i2sCfg.pin_bck        = AMP_BCLK;
    i2sCfg.pin_ws         = AMP_LRC;
    i2sCfg.pin_data       = AMP_DIN;
    i2sCfg.sample_rate    = PLAYBACK_RATE;
    i2sCfg.bits_per_sample = 16;
    i2sCfg.channels       = 1;
    i2sOut.begin(i2sCfg);
    EncodedAudioStream decoded(&i2sOut, &decoder);
    decoded.begin();
    StreamCopy copier(decoded, mp3Stream);
    while (mp3Stream.available()) copier.copy();
    decoded.end();
    i2sOut.end();
}

// ─────────────────────────────────────────────────────
// ── History management ────────────────────────────────
// ─────────────────────────────────────────────────────
void pushHistory(const String& role, const String& content) {
    if (historyCount < MAX_HISTORY) {
        history[historyCount++] = {role, content};
    } else {
        for (int i = 1; i < MAX_HISTORY - 1; i++) history[i] = history[i + 1];
        history[MAX_HISTORY - 1] = {role, content};
    }
}