#include <WiFi.h>
#include <WebSocketsClient.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ArduinoJson.h>  // 包含 ArduinoJson 库

// WiFi配置
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// WebSocket服务器配置
const char* ws_server = "frp-act.com";  // 替换为运行Python程序的电脑IP
const int ws_port = 22791;  // 替换为Python WebSocket服务器的端口

// 按键引脚
const int keyPin = 2;  // GPIO2
bool lastKeyState = HIGH;

// 屏幕配置
#define SCREEN_WIDTH 128  // 屏幕宽度（像素）
#define SCREEN_HEIGHT 64  // 屏幕高度（像素）
#define OLED_RESET -1     // OLED复位引脚，如果没有可以设置为 -1

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

WebSocketsClient webSocket;

void webSocketEvent(WStype_t type, uint8_t* payload, size_t length) {
    String message;
    DynamicJsonDocument doc(1024);

    switch (type) {
        case WStype_DISCONNECTED:
            Serial.println("WebSocket disconnected");
            break;
        case WStype_CONNECTED:
            Serial.println("WebSocket connected");
            break;
        case WStype_TEXT:
            Serial.println("Received text message from server:");
            message = String((char*)payload);
            Serial.println(message);

            // 解析JSON并更新屏幕
            deserializeJson(doc, message);
            if (doc["type"] == "display_update") {
                String content = doc["content"];
                String position = doc["position"];
                updateScreen(content, position);
            }
            break;
        default:
            Serial.println("Unknown WebSocket message type");
            break;
    }
}

void connectWebSocket() {
    webSocket.begin(ws_server, ws_port);
    webSocket.onEvent(webSocketEvent);
    webSocket.loop();
}

void setup() {
    Serial.begin(115200);
    pinMode(keyPin, INPUT_PULLUP);

    // 初始化WiFi
    WiFi.begin(ssid, password);
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println("WiFi connected");

    // 初始化WebSocket
    connectWebSocket();

    // 初始化屏幕
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {  // 0x3C 是屏幕的 I2C 地址
        Serial.println(F("SSD1306 allocation failed"));
        for (;;);
    }
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 10);
    display.println("Waiting for input...");
    display.display();
}

void loop() {
    // 检测按键状态
    bool keyState = digitalRead(keyPin);
    if (keyState == LOW && lastKeyState == HIGH) {  // 按键按下
        DynamicJsonDocument doc(256);
        doc["type"] = "key_press";
        doc["key_id"] = 1;
        doc["pressed"] = true;
        String message;
        serializeJson(doc, message);
        webSocket.sendTXT(message);
    }
    lastKeyState = keyState;

    // 保持WebSocket连接
    webSocket.loop();
    delay(20);
}

void updateScreen(String content, String position) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);

    if (position == "center") {
        display.setCursor((SCREEN_WIDTH - content.length() * 6) / 2, 30);
    } else if (position == "top") {
        display.setCursor(0, 10);
    } else if (position == "bottom") {
        display.setCursor(0, 50);
    } else {
        display.setCursor(0, 30);
    }

    display.println(content);
    display.display();
}