#include <WiFi.h>
#include <WebServer.h>
#include <Wire.h>
#include <Adafruit_PWMServoDriver.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <ArduinoOTA.h>          // For Over-the-Air updates
#include <Adafruit_Sensor.h>     // For DHT sensors
#include <DHT.h>                 // For DHT sensors

// --- Wi-Fi Configuration ---
const char* ssid = "Aussie Broadband 8729";
const char* password = "Ffdfmunfca";

// --- MQTT Configuration ---
const char* mqtt_server = "192.168.20.30";
const int   mqtt_port = 1883;
const char* mqtt_user = "mqtt-user";
const char* mqtt_password = "sam4jo";
const char* mqtt_control_topic = "homeassistant/reptile_feeder/control";
const char* mqtt_state_topic = "homeassistant/reptile_feeder/state";
const char* ota_password = "sam4jo"; // Password for OTA updates

// --- Hardware Pins ---
#define I2C_SDA_PIN 10
#define I2C_SCL_PIN 11
#define DHTPIN_1 6  // Pin for the first DHT22 sensor
#define DHTPIN_2 7  // Pin for the second DHT22 sensor
#define DHTPIN_3 18 // Pin for the third DHT22 sensor
#define DHTTYPE DHT22

// --- Timing Configuration ---
#define SENSOR_READ_INTERVAL 30000 // Time in ms between sensor reads (N). 30 seconds.

// --- Servo Configuration ---
#define SERVO_COUNT 16
#define SERVO_FREQ 50
#define SERVOMIN 150
#define SERVOMAX 600

// --- Global Objects ---
WebServer server(80);
Adafruit_PWMServoDriver pca = Adafruit_PWMServoDriver();
WiFiClient espClient;
PubSubClient mqttClient(espClient);
DHT dht_sensors[] = { DHT(DHTPIN_1, DHTTYPE), DHT(DHTPIN_2, DHTTYPE), DHT(DHTPIN_3, DHTTYPE) };
const int num_sensors = sizeof(dht_sensors) / sizeof(dht_sensors[0]);
uint8_t current_angles[SERVO_COUNT];
unsigned long previousSensorMillis = 0;

// --- Function Declarations ---
void move_servo_smoothly(uint8_t servo_num, uint8_t target_angle, uint8_t speed);
void parse_serial_command(String cmd);
void handle_web_move();
void handle_root();
void reconnectMQTT();
void mqtt_callback(char* topic, byte* payload, unsigned int length);
void readAndPublishSensors();
void setupOTA();

void setup() {
    Serial.begin(115200);
    delay(1000);
    Serial.println("\nReptile Feeder booting up...");

    // Initialize Servos
    Wire.begin(I2C_SDA_PIN, I2C_SCL_PIN);
    pca.begin();
    pca.setPWMFreq(SERVO_FREQ);
    for (uint8_t i = 0; i < SERVO_COUNT; i++) {
        current_angles[i] = 0;
        uint16_t pulselen = map(0, 0, 180, SERVOMIN, SERVOMAX);
        pca.setPWM(i, 0, pulselen);
    }
    Serial.println("Servos initialized.");

    // Initialize DHT Sensors
    for (int i = 0; i < num_sensors; i++) {
        dht_sensors[i].begin();
    }
    Serial.println("DHT sensors initialized.");

    // Connect to Wi-Fi
    Serial.printf("Connecting to %s ", ssid);
    WiFi.begin(ssid, password);
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println(" Connected!");
    Serial.print("IP Address: ");
    Serial.println(WiFi.localIP());

    // Setup MQTT
    mqttClient.setServer(mqtt_server, mqtt_port);
    mqttClient.setCallback(mqtt_callback);

    // Setup Web Server
    server.on("/", handle_root);
    server.on("/move", handle_web_move);
    server.begin();
    Serial.println("Web server started.");

    // Setup OTA Updates
    setupOTA();
    Serial.println("OTA ready. Hostname: reptile-feeder");
}

void loop() {
    ArduinoOTA.handle(); // Handle OTA update requests

    if (!mqttClient.connected()) {
        reconnectMQTT();
    }
    mqttClient.loop();

    server.handleClient();

    if (Serial.available() > 0) {
        String command = Serial.readStringUntil('\n');
        command.trim();
        parse_serial_command(command);
    }

    // Non-blocking sensor read timer
    unsigned long currentMillis = millis();
    if (currentMillis - previousSensorMillis >= SENSOR_READ_INTERVAL) {
        previousSensorMillis = currentMillis;
        readAndPublishSensors();
    }
}

// --- MQTT Functions ---
void reconnectMQTT() {
    while (!mqttClient.connected()) {
        Serial.print("Attempting MQTT connection...");
        String clientId = "ReptileFeeder-"; clientId += String(random(0xffff), HEX);
        if (mqttClient.connect(clientId.c_str(), mqtt_user, mqtt_password)) {
            Serial.println("connected");
            mqttClient.subscribe(mqtt_control_topic);
            mqttClient.publish(mqtt_state_topic, "online");
        } else {
            Serial.print("failed, rc="); Serial.print(mqttClient.state());
            Serial.println(" try again in 5 seconds");
            delay(5000);
        }
    }
}

void mqtt_callback(char* topic, byte* payload, unsigned int length) {
    Serial.print("Message arrived on topic: "); Serial.println(topic);
    JsonDocument doc;
    DeserializationError error = deserializeJson(doc, payload, length);
    if (error) { Serial.print("deserializeJson() failed: "); Serial.println(error.c_str()); return; }
    if (!doc.containsKey("servo") || !doc.containsKey("angle")) { Serial.println("Error: JSON must contain 'servo' and 'angle'"); return; }
    int servo_num = doc["servo"];
    int angle = doc["angle"];
    int speed = doc.containsKey("speed") ? doc["speed"].as<int>() : 10;
    move_servo_smoothly((uint8_t)servo_num, (uint8_t)angle, (uint8_t)speed);
}

// --- Sensor Reading Function ---
void readAndPublishSensors() {
    for (int i = 0; i < num_sensors; i++) {
        float h = dht_sensors[i].readHumidity();
        float t = dht_sensors[i].readTemperature();

        if (isnan(h) || isnan(t)) {
            Serial.printf("Failed to read from DHT sensor #%d\n", i + 1);
            continue;
        }

        JsonDocument doc;
        doc["temperature"] = t;
        doc["humidity"] = h;

        char buffer[128];
        serializeJson(doc, buffer);

        String topic = String(mqtt_state_topic) + "/sensor" + String(i + 1);
        if (mqttClient.publish(topic.c_str(), buffer)) {
            Serial.printf("Published to %s: %s\n", topic.c_str(), buffer);
        } else {
            Serial.printf("Failed to publish to %s\n", topic.c_str());
        }
    }
}

// --- OTA Setup Function ---
void setupOTA() {
    ArduinoOTA.setHostname("reptile-feeder");
    ArduinoOTA.setPassword(ota_password);
    ArduinoOTA
    .onStart([]() {
        String type;
        if (ArduinoOTA.getCommand() == U_FLASH)
            type = "sketch";
        else // U_SPIFFS
            type = "filesystem";
        Serial.println("Start updating " + type);
    })
    .onEnd([]() {
        Serial.println("\nEnd");
    })
    .onProgress([](unsigned int progress, unsigned int total) {
        Serial.printf("Progress: %u%%\r", (progress / (total / 100)));
    })
    .onError([](ota_error_t error) {
        Serial.printf("Error[%u]: ", error);
        if (error == OTA_AUTH_ERROR) Serial.println("Auth Failed");
        else if (error == OTA_BEGIN_ERROR) Serial.println("Begin Failed");
        else if (error == OTA_CONNECT_ERROR) Serial.println("Connect Failed");
        else if (error == OTA_RECEIVE_ERROR) Serial.println("Receive Failed");
        else if (error == OTA_END_ERROR) Serial.println("End Failed");
    });
    ArduinoOTA.begin();
}


// --- Web, Serial, and Servo Functions (Unchanged) ---
void handle_root() { server.send(200, "text/html", "<h1>Reptile Feeder</h1>"); }
void handle_web_move() {
    int s = server.arg("servo").toInt();
    int a = server.arg("angle").toInt();
    int v = server.arg("speed").toInt();
    move_servo_smoothly(s, a, v);
    server.send(200, "text/plain", "OK");
}
void parse_serial_command(String cmd) {
    char buf[50];
    cmd.toCharArray(buf, 50);
    int s, a, v;
    if (sscanf(buf, "move %d %d %d", &s, &a, &v) == 3) {
        move_servo_smoothly(s, a, v);
    } else {
        Serial.println("Invalid format");
    }
}
void move_servo_smoothly(uint8_t servo_num, uint8_t target_angle, uint8_t speed) {
    if (servo_num >= SERVO_COUNT || target_angle > 180 || speed < 1 || speed > 10) {
        Serial.println("Error: Move parameters out of range.");
        return;
    }
    uint8_t start_angle = current_angles[servo_num];
    if (target_angle == start_angle) return;
    uint8_t step_delay = (11 - speed) * 5;
    Serial.printf("Moving servo %d from %d to %d at speed %d\n", servo_num, start_angle, target_angle, speed);
    int step = (target_angle > start_angle) ? 1 : -1;
    for (int i = start_angle; i != (int)target_angle + step; i += step) {
        uint16_t pulselen = map(i, 0, 180, SERVOMIN, SERVOMAX);
        pca.setPWM(servo_num, 0, pulselen);
        delay(step_delay);
    }
    current_angles[servo_num] = target_angle;
    Serial.println("Move complete.");
}