#include <WiFi.h>
#include <ArduinoHA.h>

// --- Existing display / sensor includes ---
#include "OLED_Display.h"
#include <SensirionI2CSen5x.h>
#include <Wire.h>

// ── WiFi / MQTT credentials ──────────────────────────────────────────────────
const char* WIFI_SSID     = "";       // Your WIFI Name
const char* WIFI_PASSWORD = "";       // Your WIFI Password
const char* MQTT_HOST     = "";       // Home Assistant ip, e.g. "192.168.1.100"
const char* MQTT_USER     = "";       // MQTT credentials from Home Assistant
const char* MQTT_PASSWORD = "";       // MQTT credentials from Home Assistant
// ─────────────────────────────────────────────────────────────────────────────

OLED_Display Display(18, 19, 0x3C); // (SDA Pin No. , SCL Pin No. , I2C Address)
TwoWire SENS(i2c0, 4, 5); // Sensorion Air Sensor I2C connection.
SensirionI2CSen5x sen55;

WiFiClient wifiClient;
HADevice   device("pico-sen55");   // Must be unique on your network
HAMqtt     mqtt(wifiClient, device);

// Declare one HA sensor entity per reading
HASensorNumber sensorPm1p0    ("pm1p0",    HASensorNumber::PrecisionP2);
HASensorNumber sensorPm2p5    ("pm2p5",    HASensorNumber::PrecisionP2);
HASensorNumber sensorPm4p0    ("pm4p0",    HASensorNumber::PrecisionP2);
HASensorNumber sensorPm10p0   ("pm10p0",   HASensorNumber::PrecisionP2);
HASensorNumber sensorHumidity ("humidity", HASensorNumber::PrecisionP2);
HASensorNumber sensorTemp     ("temp",     HASensorNumber::PrecisionP2);
HASensorNumber sensorVoc      ("voc",      HASensorNumber::PrecisionP0);
HASensorNumber sensorNox      ("nox",      HASensorNumber::PrecisionP0);

void connectWifi() {
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    Serial.print("Connecting to WiFi");
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println(" connected.");
}

void setup() {
    Serial.begin(115200);
    delay(1000);

    Display.begin();
    SENS.begin();
    sen55.begin(SENS);

    uint16_t error;

    // Reset sensor
    error = sen55.deviceReset();
    if (error) {
        Serial.print("Error resetting device: ");
        Serial.println(error);
    }
    delay(1000);

    // Start measurement
    error = sen55.startMeasurement();
    if (error) {
        Serial.print("Error starting measurement. Error code: ");
        Serial.println(error);
    } else {
        Serial.println("Measurement started.");
    }

    // ── WiFi ─────────────────────────────────────────────────────────────────
    connectWifi();

    // ── HA device metadata ───────────────────────────────────────────────────
    device.setName("Air Quality Sensor");
    device.setModel("Pico W + SEN55");
    device.setManufacturer("DIY");

    // ── Sensor metadata (sets icons, units & device classes in HA) ───────────
    sensorPm1p0.setName("PM 1.0");
    sensorPm1p0.setUnitOfMeasurement("µg/m³");
    sensorPm1p0.setDeviceClass("pm1");

    sensorPm2p5.setName("PM 2.5");
    sensorPm2p5.setUnitOfMeasurement("µg/m³");
    sensorPm2p5.setDeviceClass("pm25");

    sensorPm4p0.setName("PM 4.0");
    sensorPm4p0.setUnitOfMeasurement("µg/m³");

    sensorPm10p0.setName("PM 10");
    sensorPm10p0.setUnitOfMeasurement("µg/m³");
    sensorPm10p0.setDeviceClass("pm10");

    sensorHumidity.setName("Humidity");
    sensorHumidity.setUnitOfMeasurement("%");
    sensorHumidity.setDeviceClass("humidity");

    sensorTemp.setName("Temperature");
    sensorTemp.setUnitOfMeasurement("°C");
    sensorTemp.setDeviceClass("temperature");

    sensorVoc.setName("VOC Index");
    sensorVoc.setUnitOfMeasurement("index");

    sensorNox.setName("NOx Index");
    sensorNox.setUnitOfMeasurement("index");

    // ── Connect to MQTT broker ────────────────────────────────────────────────
    mqtt.begin(MQTT_HOST, MQTT_USER, MQTT_PASSWORD);
}

void loop() {
    mqtt.loop();   // maintains MQTT connection & handles discovery

    uint16_t error;

    float massConcentrationPm1p0;
    float massConcentrationPm2p5;
    float massConcentrationPm4p0;
    float massConcentrationPm10p0;
    float ambientHumidity;
    float ambientTemperature;
    float vocIndex;
    float noxIndex;

    error = sen55.readMeasuredValues(
        massConcentrationPm1p0, massConcentrationPm2p5,
        massConcentrationPm4p0, massConcentrationPm10p0,
        ambientHumidity, ambientTemperature, vocIndex, noxIndex);

    if (error) {
        Serial.print("Error code: ");
        Serial.println(error);
    } else {
        // ── Update OLED (unchanged logic) ────────────────────────────────────
        char buffer[10];
        sprintf(buffer, "%.2f ug/m   ", massConcentrationPm1p0);  Display.writeLine(buffer, 1, 9 * 5 - 1);
        sprintf(buffer, "%.2f ug/m   ", massConcentrationPm2p5);  Display.writeLine(buffer, 2, 9 * 5 - 1);
        sprintf(buffer, "%.2f ug/m   ", massConcentrationPm10p0); Display.writeLine(buffer, 3, 9 * 5 - 1);
        sprintf(buffer, "%.2f %RH   ", ambientHumidity);          Display.writeLine(buffer, 4, 12 * 5 - 1);
        sprintf(buffer, "%.2f C      ", ambientTemperature);      Display.writeLine(buffer, 5, 7 * 5);
        sprintf(buffer, "%.0f   ", vocIndex);                     Display.writeLine(buffer, 6, 13 * 5);
        sprintf(buffer, "%.0f   ", noxIndex);                     Display.writeLine(buffer, 7, 13 * 5);

        // ── Publish to Home Assistant every loop (1 s interval from delay) ───
        sensorPm1p0.setValue(massConcentrationPm1p0);
        sensorPm2p5.setValue(massConcentrationPm2p5);
        sensorPm4p0.setValue(massConcentrationPm4p0);
        sensorPm10p0.setValue(massConcentrationPm10p0);
        sensorHumidity.setValue(ambientHumidity);
        sensorTemp.setValue(ambientTemperature);
        sensorVoc.setValue(vocIndex);
        sensorNox.setValue(noxIndex);
    }

    Display.writeLine("---  Air Sensor!  ---", 0, 0);
    Display.writeLine("PM 1.0:", 1, 0);
    Display.writeLine("PM 2.5:", 2, 0);
    Display.writeLine("PM 10 :", 3, 0);
    Display.writeLine("Humidity:", 4, 0);
    Display.writeLine("Temp:", 5, 0);
    Display.writeLine("VOC Index:", 6, 0);
    Display.writeLine("NOx Index:", 7, 0);

    delay(1000); // 1 second between reads
}