#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <DHT.h>

// === CONFIG ===
#define I2C_SCL_LCD 22    // LCD I2C pins
#define I2C_SDA_LCD 21

#define I2C_SCL_CCS 19    // CCS811 I2C pins  
#define I2C_SDA_CCS 18

#define LCD_ADDR 0x27
#define LCD_COLS 20
#define LCD_ROWS 4

#define BUTTON_PIN 13
#define BUZZER_PIN 33
#define LED_PINS {12, 14, 27, 26, 25}

#define DHT_PIN 32
#define DHT_TYPE DHT22
#define DUST_LED_PIN 15
#define DUST_VO_PIN 34

// CCS811 Config
#define CCS811_ADDR 0x5A

// WiFi + MQTT
const char* WIFI_SSID = "Wokwi-GUEST";
const char* WIFI_PASS = "";
const char* MQTT_SERVER = "broker.hivemq.com";
const int MQTT_PORT = 1883;
const char* CLIENT_ID = "ESP32-Arduino";
const char* TOPIC_DATA = "23127007_23127017_23127236/ESP_DATA";
const char* TOPIC_CONTROL = "23127007_23127017_23127236/WEB_DATA";

// Global variables
StaticJsonDocument<512> receivedData;
bool buzzerActive = false;
bool lastButtonState = true;
int qualityLevel = 0;
bool dataReceived = false;

// Hardware objects
LiquidCrystal_I2C lcd(LCD_ADDR, LCD_COLS, LCD_ROWS);
DHT dht(DHT_PIN, DHT_TYPE);
WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);

// LED pins array
int ledPins[] = LED_PINS;
const int numLeds = sizeof(ledPins) / sizeof(ledPins[0]);

// === CCS811 Class ===
class CCS811 {
private:
    // Register addresses
    static const uint8_t STATUS = 0x00;
    static const uint8_t MEAS_MODE = 0x01;
    static const uint8_t ALG_RESULT_DATA = 0x02;
    static const uint8_t HW_ID = 0x20;
    static const uint8_t APP_START = 0xF4;
    static const uint8_t ENV_DATA = 0x05;
    
    TwoWire* i2c;
    uint8_t addr;
    uint16_t co2;
    uint16_t tvoc;
    bool initialized;
    
public:
    CCS811(TwoWire* wire, uint8_t address = 0x5A) {
        i2c = wire;
        addr = address;
        co2 = 400;
        tvoc = 0;
        initialized = false;
    }
    
    bool begin() {
        try {
            Serial.println("CCS811: Starting init...");
            
            // Simple direct approach like MicroPython
            i2c->beginTransmission(addr);
            i2c->write(HW_ID);
            uint8_t result = i2c->endTransmission(false); // Repeated start
            Serial.printf("Write HW_ID register result: %d\n", result);
            
            uint8_t bytesReceived = i2c->requestFrom(addr, (uint8_t)1);
            Serial.printf("Bytes requested: 1, received: %d\n", bytesReceived);
            
            if (i2c->available()) {
                uint8_t hwId = i2c->read();
                Serial.printf("CCS811: Read HW ID: 0x%02x (expected 0x81)\n", hwId);
                if (hwId != 0x81) {
                    Serial.printf("CCS811: Wrong HW ID: 0x%02x\n", hwId);
                    return false;
                }
            } else {
                Serial.println("CCS811: No data available after requestFrom");
                return false;
            }
            
            // Start app
            Serial.println("CCS811: Starting app...");
            i2c->beginTransmission(addr);
            i2c->write(APP_START);
            result = i2c->endTransmission();
            Serial.printf("App start result: %d\n", result);
            if (result != 0) return false;
            delay(100);
            
            // Set measurement mode
            Serial.println("CCS811: Setting measurement mode...");
            i2c->beginTransmission(addr);
            i2c->write(MEAS_MODE);
            i2c->write(0x10);
            result = i2c->endTransmission();
            Serial.printf("Meas mode result: %d\n", result);
            if (result != 0) return false;
            delay(100);
            
            initialized = true;
            Serial.println("CCS811: Initialized successfully");
            return true;
            
        } catch (...) {
            Serial.println("CCS811: Init failed");
            return false;
        }
    }
    
private:
    // Helper function equivalent to MicroPython's readfrom_mem()
    bool readRegister(uint8_t reg, uint8_t* data, uint8_t len) {
        // Write register address
        i2c->beginTransmission(addr);
        i2c->write(reg);
        if (i2c->endTransmission(false) != 0) return false;  // Send repeated start
        
        // Read data
        if (i2c->requestFrom(addr, len) != len) return false;
        
        for (uint8_t i = 0; i < len; i++) {
            if (i2c->available()) {
                data[i] = i2c->read();
            } else {
                return false;
            }
        }
        return true;
    }
    
    // Helper function equivalent to MicroPython's writeto()
    bool writeCommand(uint8_t cmd) {
        i2c->beginTransmission(addr);
        i2c->write(cmd);
        return (i2c->endTransmission() == 0);
    }
    
    // Helper function equivalent to MicroPython's writeto_mem()
    bool writeRegister(uint8_t reg, uint8_t* data, uint8_t len) {
        i2c->beginTransmission(addr);
        i2c->write(reg);
        for (uint8_t i = 0; i < len; i++) {
            i2c->write(data[i]);
        }
        return (i2c->endTransmission() == 0);
    }

public:
    
    bool dataAvailable() {
        if (!initialized) return false;
        
        uint8_t status;
        if (!readRegister(STATUS, &status, 1)) return false;
        return (status & 0x08) != 0;  // Data ready bit
    }
    
    bool readData() {
        if (!initialized) return false;
        
        try {
            uint8_t data[4];
            if (!readRegister(ALG_RESULT_DATA, data, 4)) {
                Serial.println("CCS811: Read error");
                return false;
            }
            
            co2 = (data[0] << 8) | data[1];
            tvoc = (data[2] << 8) | data[3];
            return true;
            
        } catch (...) {
            Serial.println("CCS811: Read error");
            return false;
        }
    }
    
    uint16_t getCO2() { return co2; }
    uint16_t getTVOC() { return tvoc; }
    
    void setEnvironmentalData(float humidity, float temperature) {
        if (!initialized) return;
        
        try {
            // Convert to CCS811 format
            uint16_t humRaw = (uint16_t)(humidity * 512);
            uint16_t tempRaw = (uint16_t)((temperature + 25) * 512);
            
            uint8_t envData[4] = {
                (humRaw >> 8) & 0xFF,
                humRaw & 0xFF,
                (tempRaw >> 8) & 0xFF,
                tempRaw & 0xFF
            };
            
            writeRegister(ENV_DATA, envData, 4);
        } catch (...) {
            Serial.println("CCS811: Env data error");
        }
    }
};

// Initialize CCS811 on separate I2C bus
TwoWire I2C_CCS = TwoWire(1);
CCS811 ccs811(&I2C_CCS, CCS811_ADDR);
bool ccs811Ok = false;

// === FUNCTIONS ===
float readDust() {
    int raw = analogRead(DUST_VO_PIN);
    float voltage = raw * 3.3 / 4095.0;
    
    const float VOC_TYPICAL = 0.9;
    const float SENSITIVITY_V_PER_100UG = 0.5;
    
    float deltaV = voltage - VOC_TYPICAL;
    
    if (deltaV <= 0) {
        return 0.0;
    } else {
        return (deltaV / SENSITIVITY_V_PER_100UG) * 100;
    }
}

void readAirQuality(uint16_t* co2, uint16_t* tvoc) {
    if (!ccs811Ok) {
        *co2 = 400;
        *tvoc = 0;
        return;
    }
    
    if (ccs811.dataAvailable()) {
        if (ccs811.readData()) {
            *co2 = ccs811.getCO2();
            *tvoc = ccs811.getTVOC();
            return;
        }
    }
    
    *co2 = ccs811.getCO2();
    *tvoc = ccs811.getTVOC();
}

void wifiConnect() {
    WiFi.begin(WIFI_SSID, WIFI_PASS);
    while (WiFi.status() != WL_CONNECTED) {
        Serial.println("Connecting WiFi...");
        delay(1000);
    }
    Serial.print("WiFi connected: ");
    Serial.println(WiFi.localIP());
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
    Serial.printf("Received: %s\n", topic);
    
    if (strcmp(topic, TOPIC_CONTROL) == 0) {
        // Parse JSON payload
        StaticJsonDocument<512> doc;
        deserializeJson(doc, payload, length);
        
        receivedData = doc;
        dataReceived = true;
        qualityLevel = doc["quality"] | 0;
        
        Serial.printf("Quality level: %d\n", qualityLevel);
        
        // Control LEDs based on quality level
        controlLeds(qualityLevel);
        
        // Control buzzer for poor air quality
        if (qualityLevel < 2) {
            buzzerActive = true;
            activateBuzzer();
        } else {
            buzzerActive = false;
            deactivateBuzzer();
        }
    }
}

void controlLeds(int quality) {
    int ledsToTurnOn = quality + 1;
    
    // Turn off all LEDs first
    for (int i = 0; i < numLeds; i++) {
        digitalWrite(ledPins[i], LOW);
    }
    
    // Turn on the required number of LEDs
    for (int i = 0; i < min(ledsToTurnOn, numLeds); i++) {
        digitalWrite(ledPins[i], HIGH);
    }
}

void activateBuzzer() {
    ledcAttach(BUZZER_PIN, 1000, 8);  // Pin, frequency, resolution
    ledcWrite(BUZZER_PIN, 128);  // 50% duty cycle
}

void deactivateBuzzer() {
    ledcWrite(BUZZER_PIN, 0);
}

void checkButton() {
    bool currentButtonState = digitalRead(BUTTON_PIN);
    
    // Detect button press (falling edge)
    if (lastButtonState && !currentButtonState) {
        if (buzzerActive) {
            buzzerActive = false;
            deactivateBuzzer();
            Serial.println("Button pressed - buzzer turned off");
        }
    }
    
    lastButtonState = currentButtonState;
}

void mqttConnect() {
    mqttClient.setServer(MQTT_SERVER, MQTT_PORT);
    mqttClient.setCallback(mqttCallback);
    
    while (!mqttClient.connected()) {
        Serial.println("Connecting to MQTT...");
        if (mqttClient.connect(CLIENT_ID)) {
            Serial.println("Connected to MQTT broker");
            mqttClient.subscribe(TOPIC_CONTROL);
            Serial.printf("Subscribed to %s\n", TOPIC_CONTROL);
        } else {
            Serial.printf("MQTT connection failed, rc=%d\n", mqttClient.state());
            delay(5000);
        }
    }
}

void displayData() {
    lcd.clear();
    
    if (dataReceived) {
        // Line 1: Temperature & Humidity
        lcd.setCursor(0, 0);
        float temp = receivedData["temperature"] | 0.0;
        int hum = receivedData["humidity"] | 0;
        lcd.printf("T:%.1fC H:%d%%", temp, hum);
        
        // Line 2: Dust values
        lcd.setCursor(0, 1);
        int pm25 = receivedData["pm25"] | 0;
        int pm10 = receivedData["pm10"] | 0;
        lcd.printf("PM2.5:%d PM10:%d", pm25, pm10);
        
        // Line 3: Air quality
        lcd.setCursor(0, 2);
        int co2 = receivedData["co2"] | 0;
        int voc = receivedData["voc"] | 0;
        lcd.printf("CO2:%d VOC:%d", co2, voc);
        
        // Line 4: Quality level
        lcd.setCursor(0, 3);
        const char* qualityTexts[] = {"Very Poor", "Poor", "Fair", "Good", "Excellent"};
        lcd.printf("Quality: %s", qualityTexts[constrain(qualityLevel, 0, 4)]);
        
    } else {
        // No data received yet - show sensor readings
        float temp = dht.readTemperature();
        float hum = dht.readHumidity();
        float dust = readDust();
        uint16_t co2, tvoc;
        readAirQuality(&co2, &tvoc);
        
        if (!isnan(temp) && !isnan(hum)) {
            lcd.setCursor(0, 0);
            lcd.printf("T:%.1fC H:%.0f%%", temp, hum);
            lcd.setCursor(0, 1);
            lcd.printf("Dust:%.0f ug/m3", dust);
            lcd.setCursor(0, 2);
            lcd.printf("CO2:%d TVOC:%d", co2, tvoc);
            lcd.setCursor(0, 3);
            lcd.print("Waiting for data...");
        } else {
            lcd.setCursor(0, 0);
            lcd.print("Sensor Error");
        }
    }
}

// === SETUP ===
void setup() {
    Serial.begin(115200);
    
    // Initialize I2C buses
    Wire.begin(I2C_SDA_LCD, I2C_SCL_LCD);  // Default I2C for LCD
    I2C_CCS.begin(I2C_SDA_CCS, I2C_SCL_CCS, 400000);  // Separate I2C for CCS811
    
    // Initialize LCD
    lcd.init();
    lcd.backlight();
    
    // Initialize CCS811
    ccs811Ok = ccs811.begin();
    
    // Initialize pins
    pinMode(BUTTON_PIN, INPUT_PULLUP);
    for (int i = 0; i < numLeds; i++) {
        pinMode(ledPins[i], OUTPUT);
        digitalWrite(ledPins[i], LOW);
    }
    
    pinMode(DUST_LED_PIN, OUTPUT);
    digitalWrite(DUST_LED_PIN, HIGH);
    
    // Initialize DHT
    dht.begin();
    
    // Connect to WiFi and MQTT
    wifiConnect();
    mqttConnect();
    
    Serial.printf("CCS811 Status: %s\n", ccs811Ok ? "OK" : "FAILED");
}

float add_noise() {
  float r = static_cast <float> (rand()) / static_cast <float> (RAND_MAX);
  return r;
}

// === MAIN LOOP ===
unsigned long lastPublish = 0;
void loop() {
    // Handle MQTT
    if (!mqttClient.connected()) {
        Serial.printf("MQTT disconnected\n");
        mqttConnect();
    }
    mqttClient.loop();
    
    // Check button for buzzer control
    checkButton();
    
    unsigned long currentTime = millis();
    
    // Publish sensor data every 10 seconds
    if (currentTime - lastPublish >= 10000) {
        // Read DHT22
        uint16_t temp = round(dht.readTemperature() + add_noise());
        uint16_t hum = round(dht.readHumidity() + add_noise());
        
        // Read dust sensor
        uint16_t dust = round(readDust() + add_noise());
        
        // Read CCS811 air quality
        uint16_t co2, tvoc;
        readAirQuality(&co2, &tvoc);
        co2 = (uint16_t)(co2 + add_noise());
        tvoc = (uint16_t)(tvoc + add_noise());
        
        Serial.printf("Readings: T:%.1f H:%.1f Dust:%.1f CO2:%d TVOC:%d\n", 
                     temp, hum, dust, co2, tvoc);
        
        // Send environmental data to CCS811 for compensation
        if (ccs811Ok && !isnan(temp) && !isnan(hum)) {
            ccs811.setEnvironmentalData(hum, temp);
        }
        
        // Build JSON payload
        StaticJsonDocument<256> doc;
        doc["temperature"] = temp;
        doc["humidity"] = hum;
        doc["pm25"] = dust;
        doc["pm10"] = 0;
        doc["co2"] = co2;
        doc["voc"] = tvoc;
        
        String jsonString;
        serializeJson(doc, jsonString);
        
        Serial.println("Publishing sensor data: " + jsonString);
        mqttClient.publish(TOPIC_DATA, jsonString.c_str());
        
        lastPublish = currentTime;
    }
    
    // Update display
    displayData();
    
    delay(1000);
}