#include <WiFi.h>
#include <PubSubClient.h>
#include <ESP32Servo.h>
#include "DHTesp.h"
#include <math.h>
#include <Preferences.h>

#define LDR_PIN 32
#define SERVO_PIN 26
#define DHT_PIN 14

WiFiClient espClient;
PubSubClient mqttClient(espClient);
Servo servo;
DHTesp dhtSensor;
Preferences preferences;

const char* ssid = "Wokwi-GUEST";
const char* password = "";

const char* mqttServer = "broker.hivemq.com";
const int mqttPort = 1883;

const char* LIGHT_AVG_TOPIC = "MEDIBOX/LIGHT_AVG";
const char* SERVO_ANGLE_TOPIC = "MEDIBOX/SERVO_ANGLE";
const char* MIN_ANGLE_TOPIC = "MEDIBOX/MIN_ANGLE";
const char* CONTROL_FACTOR_TOPIC = "MEDIBOX/CONTROL_FACTOR";
const char* TMED_TOPIC = "MEDIBOX/TMED";
const char* SAMPLING_INTERVAL_TOPIC = "MEDIBOX/SAMPLING_INTERVAL";
const char* SENDING_INTERVAL_TOPIC = "MEDIBOX/SENDING_INTERVAL";

int minAngle = 30;           // θoffset (default: 30 degrees)
float controlFactor = 0.75;  // γ (default: 0.75)
float Tmed = 30.0;           // Ideal temperature (default: 30°C)
int samplingInterval = 5;    // ts (default: 5 seconds)
int sendingInterval = 120;   // tu (default: 120 seconds)

#define MAX_SAMPLES 60
float lightSamples[MAX_SAMPLES];
volatile int sampleCount = 0;

float currentLightIntensity = 0.5;
float currentTemperature = 25.0;

unsigned long lastSampleTime = 0;
unsigned long lastSendTime = 0;
unsigned long lastTemperatureReadTime = 0;
unsigned long lastMqttReconnectAttempt = 0;

bool servoInitialized = false;
int servoInitAttempts = 0;

void setupWifi();
void connectToBroker();
void receiveCallback(char* topic, byte* payload, unsigned int length);
void publishCurrentConfig();
void readTemperature();
void sampleLightIntensity();
void updateServoPosition();
void sendAverageAndStatus();
void loadPreferences();
void savePreferences();
void initializeServo();

void setup() {
  Serial.begin(115200);
  delay(1000);
  Serial.println("\n\n--- MEDIBOX Starting ---");

  preferences.begin("medibox", false);
  loadPreferences();

  setupWifi();

  mqttClient.setServer(mqttServer, mqttPort);
  mqttClient.setCallback(receiveCallback);
  mqttClient.setKeepAlive(60);
  mqttClient.setSocketTimeout(10);

  pinMode(SERVO_PIN, OUTPUT);
  delay(1000);

  pinMode(LDR_PIN, INPUT);
  pinMode(DHT_PIN, INPUT);

  delay(1000);

  Serial.println("Initializing servo motor...");
  for (int i = 0; i < 3; i++) {
    initializeServo();
    if (servoInitialized) {
      Serial.println("Servo initialization successful on attempt " + String(i + 1));
      break;
    }
    Serial.println("Waiting before retry...");
    delay(1000);
  }

  dhtSensor.setup(DHT_PIN, DHTesp::DHT22);

  Serial.println("Medibox system initialized");
  Serial.println("Initial configuration:");
  Serial.println("--------------------");
  Serial.println("Minimum angle (θoffset): " + String(minAngle));
  Serial.println("Controlling factor (γ): " + String(controlFactor));
  Serial.println("Ideal temperature (Tmed): " + String(Tmed));
  Serial.println("Sampling interval (ts): " + String(samplingInterval) + " seconds");
  Serial.println("Sending interval (tu): " + String(sendingInterval) + " seconds");
  Serial.println("--------------------");

  if (servoInitialized) {
    servo.write(minAngle);
    delay(1000);
    Serial.println("Servo test: Setting to 90 degrees");
    servo.write(90);
    delay(1000);
    Serial.println("Servo test: Returning to " + String(minAngle));
    servo.write(minAngle);
    delay(1000);
  } else {
    Serial.println("WARNING: Servo not initialized yet, will try again during operation");
  }

  lastSampleTime = millis();
  lastSendTime = millis();
  lastTemperatureReadTime = millis();
}

void initializeServo() {
  Serial.println("Starting servo initialization sequence...");
  if (servo.attached()) {
    Serial.println("Detaching previously attached servo");
    servo.detach();
    delay(300);
  }
  pinMode(SERVO_PIN, OUTPUT);
  digitalWrite(SERVO_PIN, LOW);
  delay(200);
  Serial.println("Setting up servo parameters...");
  servo.setPeriodHertz(50);
  Serial.println("Attempting to attach servo on pin " + String(SERVO_PIN));
  int attachResult = servo.attach(SERVO_PIN, 500, 2500);
  Serial.println("Attach result: " + String(attachResult));
  servoInitialized = (attachResult != -1);
  if (servoInitialized) {
    Serial.println("Servo initialized successfully!");
    Serial.println("Attached status after initialization: " + String(servo.attached() ? "yes" : "no"));
    servo.write(minAngle);
    delay(1000);
    int pos = servo.read();
    Serial.println("Servo set to initial position: " + String(pos));
  } else {
    Serial.println("FAILED to initialize servo!");
    Serial.println("Check: 1) SERVO_PIN (" + String(SERVO_PIN) + ") is PWM-capable");
    Serial.println("      2) Servo VCC is connected to 5V (not 3.3V) in Wokwi");
    Serial.println("      3) GND is shared between ESP32 and servo");
    Serial.println("      4) Servo is physically connected in Wokwi");
    Serial.println("      5) Try pins 27 or 33 if pin 26 fails");
  }
}

void loop() {
  unsigned long currentMillis = millis();
  static unsigned long lastHeartbeat = 0;
  static unsigned long lastServoRetryTime = 0;

  if (!servoInitialized && servoInitAttempts < 3 && 
      (currentMillis - lastServoRetryTime > 15000 || currentMillis < lastServoRetryTime)) {
    lastServoRetryTime = currentMillis;
    servoInitAttempts++;
    Serial.println("Retrying servo initialization (attempt " + String(servoInitAttempts) + ")");
    initializeServo();
    if (servoInitialized) {
      servo.write(minAngle);
      delay(1000);
      Serial.println("Servo position set to " + String(minAngle) + " after successful retry");
    }
  }

  if (currentMillis - lastHeartbeat > 5000 || currentMillis < lastHeartbeat) {
    lastHeartbeat = currentMillis;
    Serial.println("Medibox running... " + String(currentMillis / 1000) + "s");
    Serial.println("Current parameters: minAngle=" + String(minAngle) + 
                   ", controlFactor=" + String(controlFactor) + 
                   ", Tmed=" + String(Tmed) +
                   ", samplingInterval=" + String(samplingInterval) +
                   ", sendingInterval=" + String(sendingInterval));
    Serial.println("Servo status: " + String(servoInitialized ? "initialized" : "not initialized") + 
                   ", attempts: " + String(servoInitAttempts) + 
                   ", attached: " + String(servo.attached() ? "yes" : "no"));
  }

  if (!mqttClient.connected() && 
      (currentMillis - lastMqttReconnectAttempt > 5000 || currentMillis < lastMqttReconnectAttempt)) {
    lastMqttReconnectAttempt = currentMillis;
    connectToBroker();
  }

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

  if (currentMillis - lastTemperatureReadTime >= 2000 || currentMillis < lastTemperatureReadTime) {
    lastTemperatureReadTime = currentMillis;
    readTemperature();
  }

  if (currentMillis - lastSampleTime >= (unsigned long)samplingInterval * 1000UL || 
      currentMillis < lastSampleTime) {
    lastSampleTime = currentMillis;
    sampleLightIntensity();
    updateServoPosition();
  }

  if (currentMillis - lastSendTime >= (unsigned long)sendingInterval * 1000UL || 
      currentMillis < lastSendTime) {
    lastSendTime = currentMillis;
    sendAverageAndStatus();
  }

  delay(50);
}

void loadPreferences() {
  if (minAngle == 30) {
    minAngle = preferences.getInt("minAngle", 30);
  }
  if (fabs(controlFactor - 0.75) < 0.001) {
    controlFactor = preferences.getFloat("ctrlFactor", 0.75);
  }
  if (fabs(Tmed - 30.0) < 0.001) {
    Tmed = preferences.getFloat("tmed", 30.0);
  }
  if (samplingInterval == 5) {
    samplingInterval = preferences.getInt("sampInt", 5);
  }
  if (sendingInterval == 120) {
    sendingInterval = preferences.getInt("sendInt", 120);
  }
  Serial.println("Loaded preferences:");
  Serial.println("minAngle: " + String(minAngle));
  Serial.println("controlFactor: " + String(controlFactor));
  Serial.println("Tmed: " + String(Tmed));
  Serial.println("samplingInterval: " + String(samplingInterval));
  Serial.println("sendingInterval: " + String(sendingInterval));
}

void savePreferences() {
  preferences.putInt("minAngle", minAngle);
  preferences.putFloat("ctrlFactor", controlFactor);
  preferences.putFloat("tmed", Tmed);
  preferences.putInt("sampInt", samplingInterval);
  preferences.putInt("sendInt", sendingInterval);
  Serial.println("Preferences saved successfully");
}

void setupWifi() {
  delay(10);
  Serial.println();
  Serial.print("Connecting to WiFi network: ");
  Serial.println(ssid);
  WiFi.begin(ssid, password);
  unsigned long startAttemptTime = millis();
  while (WiFi.status() != WL_CONNECTED && millis() - startAttemptTime < 10000) {
    delay(500);
    Serial.print(".");
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("");
    Serial.println("WiFi connected");
    Serial.print("IP address: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println("");
    Serial.println("WiFi connection failed. For Wokwi simulation, this is normal.");
    Serial.println("Will proceed without WiFi connection.");
  }
}

void connectToBroker() {
  Serial.print("Connecting to MQTT broker...");
  String clientId = "ESP32-MEDIBOX-";
  clientId += String(random(0xffff), HEX);
  bool connected = false;
  try {
    connected = mqttClient.connect(clientId.c_str());
  } catch (...) {
    Serial.println("MQTT connect threw exception");
  }
  if (connected) {
    Serial.println("connected");
    try {
      mqttClient.subscribe(MIN_ANGLE_TOPIC);
      mqttClient.subscribe(CONTROL_FACTOR_TOPIC);
      mqttClient.subscribe(TMED_TOPIC);
      mqttClient.subscribe(SAMPLING_INTERVAL_TOPIC);
      mqttClient.subscribe(SENDING_INTERVAL_TOPIC);
      Serial.println("Subscribed to control topics");
      mqttClient.publish("MEDIBOX/STATUS", "Device connected");
      publishCurrentConfig();
    } catch (...) {
      Serial.println("Exception in MQTT subscription");
    }
  } else {
    Serial.print("failed, rc=");
    Serial.println(mqttClient.state());
  }
}

void receiveCallback(char* topic, byte* payload, unsigned int length) {
  if (length > 20) length = 20;
  char message[21];
  for (int i = 0; i < length; i++) {
    message[i] = (char)payload[i];
  }
  message[length] = '\0';
  Serial.print("Message arrived [");
  Serial.print(topic);
  Serial.print("] ");
  Serial.println(message);

  // Ignore messages from CURRENT topics to prevent feedback loop
  if (strstr(topic, "MEDIBOX/CURRENT/") != NULL) {
    Serial.println("Ignoring CURRENT topic message to avoid feedback loop");
    return;
  }

  bool paramChanged = false;
  if (strcmp(topic, MIN_ANGLE_TOPIC) == 0) {
    int newValue = atoi(message);
    if (newValue >= 0 && newValue <= 90) {
      if (newValue != minAngle) {
        minAngle = newValue;
        paramChanged = true;
        savePreferences();
        Serial.println("Updated minimum angle: " + String(minAngle));
        mqttClient.publish("MEDIBOX/CURRENT/MIN_ANGLE", message, true);
      }
    }
  } else if (strcmp(topic, CONTROL_FACTOR_TOPIC) == 0) {
    float newValue = atof(message);
    if (newValue >= 0.1 && newValue <= 2.0) {
      if (fabs(newValue - controlFactor) > 0.001) {
        controlFactor = newValue;
        paramChanged = true;
        savePreferences();
        Serial.println("Updated controlling factor: " + String(controlFactor));
        mqttClient.publish("MEDIBOX/CURRENT/CONTROL_FACTOR", message, true);
      }
    }
  } else if (strcmp(topic, TMED_TOPIC) == 0) {
    float newValue = atof(message);
    if (newValue >= 0 && newValue <= 50) {
      if (fabs(newValue - Tmed) > 0.001) {
        Tmed = newValue;
        paramChanged = true;
        savePreferences();
        Serial.println("Updated ideal temperature: " + String(Tmed));
        mqttClient.publish("MEDIBOX/CURRENT/TMED", message, true);
      }
    }
  } else if (strcmp(topic, SAMPLING_INTERVAL_TOPIC) == 0) {
    int newValue = atoi(message);
    if (newValue >= 1 && newValue <= 600) {
      if (newValue != samplingInterval) {
        samplingInterval = newValue;
        paramChanged = true;
        savePreferences();
        Serial.println("Updated sampling interval: " + String(samplingInterval));
        mqttClient.publish("MEDIBOX/CURRENT/SAMPLING_INTERVAL", message, true);
      }
    }
  } else if (strcmp(topic, SENDING_INTERVAL_TOPIC) == 0) {
    int newValue = atoi(message);
    if (newValue >= 10 && newValue <= 3600) {
      if (newValue != sendingInterval) {
        sendingInterval = newValue;
        paramChanged = true;
        savePreferences();
        Serial.println("Updated sending interval: " + String(sendingInterval));
        mqttClient.publish("MEDIBOX/CURRENT/SENDING_INTERVAL", message, true);
      }
    }
  }

  if (paramChanged) {
    readTemperature();
    updateServoPosition();
    publishCurrentConfig();
  }
}

void publishCurrentConfig() {
  char msg[20];
  snprintf(msg, sizeof(msg), "%d", minAngle);
  mqttClient.publish("MEDIBOX/CURRENT/MIN_ANGLE", msg, true);
  snprintf(msg, sizeof(msg), "%.2f", controlFactor);
  mqttClient.publish("MEDIBOX/CURRENT/CONTROL_FACTOR", msg, true);
  snprintf(msg, sizeof(msg), "%.1f", Tmed);
  mqttClient.publish("MEDIBOX/CURRENT/TMED", msg, true);
  snprintf(msg, sizeof(msg), "%d", samplingInterval);
  mqttClient.publish("MEDIBOX/CURRENT/SAMPLING_INTERVAL", msg, true);
  snprintf(msg, sizeof(msg), "%d", sendingInterval);
  mqttClient.publish("MEDIBOX/CURRENT/SENDING_INTERVAL", msg, true);
  mqttClient.publish("MEDIBOX/CONFIG_STATUS", "Configuration published", true);
  Serial.println("Published current configuration");
}

void readTemperature() {
  TempAndHumidity data;
  try {
    data = dhtSensor.getTempAndHumidity();
    if (isnan(data.temperature)) {
      Serial.println("Failed to read valid temperature");
      return;
    }
    if (data.temperature < -10 || data.temperature > 60) {
      Serial.println("Temperature out of expected range");
      return;
    }
    currentTemperature = data.temperature;
    char tempStr[8];
    snprintf(tempStr, sizeof(tempStr), "%.1f", currentTemperature);
    mqttClient.publish("MEDIBOX/TEMPERATURE", tempStr);
    Serial.print("Temperature: ");
    Serial.println(currentTemperature);
  } catch (...) {
    Serial.println("Exception in temperature reading");
  }
}

void sampleLightIntensity() {
  try {
    const int numReadings = 3;
    long rawTotal = 0;
    for (int i = 0; i < numReadings; i++) {
      int reading = analogRead(LDR_PIN);
      if (reading < 0) reading = 0;
      if (reading > 4095) reading = 4095;
      rawTotal += reading;
      delay(2);
    }
    int rawValue = rawTotal / numReadings;
    float intensity = map(rawValue, 0, 4095, 100, 0) / 100.0;  // 0-1 range
    intensity = constrain(intensity, 0.0, 1.0);
    currentLightIntensity = intensity;
    if (sampleCount < MAX_SAMPLES) {
      lightSamples[sampleCount++] = intensity;
    } else {
      for (int i = 0; i < MAX_SAMPLES / 2; i++) {
        lightSamples[i] = lightSamples[i + MAX_SAMPLES / 2];
      }
      sampleCount = MAX_SAMPLES / 2;
      lightSamples[sampleCount++] = intensity;
    }
    Serial.print("Light sample: ");
    Serial.print(rawValue);
    Serial.print(" -> ");
    Serial.println(intensity, 2);
    char intStr[8];
    snprintf(intStr, sizeof(intStr), "%.2f", intensity);
    mqttClient.publish("MEDIBOX/LIGHT_INTENSITY", intStr);
  } catch (...) {
    Serial.println("Exception in light sampling");
  }
}

void updateServoPosition() {
  if (!servoInitialized) {
    Serial.println("Servo not initialized - attempting initialization");
    initializeServo();
    if (!servoInitialized) {
      Serial.println("Servo initialization failed - cannot update position");
      return;
    }
  }
  try {
    float light = constrain(currentLightIntensity, 0.0, 1.0);
    float temp = constrain(currentTemperature, 0.0, 50.0);
    int min_angle = constrain(minAngle, 0, 90);
    float factor = constrain(controlFactor, 0.1, 2.0);
    int ts = constrain(samplingInterval, 1, 600);
    int tu = constrain(sendingInterval, 10, 3600);
    float t_med = constrain(Tmed, 1.0, 50.0);

    float logTerm = (ts > 0 && tu > 0) ? fabs(log((float)ts / (float)tu)) : 0;
    float tempDeviation = (t_med > 0) ? 1 + fabs(temp - t_med) / t_med : 1.0;
    Serial.println("Servo calculation: I=" + String(light) + ", T=" + String(temp) + 
                   ", θ_offset=" + String(min_angle) + ", γ=" + String(factor) + 
                   ", ln(|ts/tu|)=" + String(logTerm) + ", |T-Tmed|/Tmed+1=" + String(tempDeviation));

    float angle = min_angle + (180 - min_angle) * light * factor * logTerm * tempDeviation;
    angle = constrain(angle, min_angle, 180);
    Serial.println("Setting servo angle: " + String(angle));

    servo.write((int)angle);
    delay(300);

    int actualAngle = servo.read();
    if (actualAngle == -1) {
      Serial.println("Warning: Servo position read failed, assuming write succeeded");
      actualAngle = (int)angle;
    }
    Serial.println("Servo position: requested=" + String((int)angle) + 
                   ", actual=" + String(actualAngle));

    char angleStr[8];
    snprintf(angleStr, sizeof(angleStr), "%.1f", angle);
    mqttClient.publish(SERVO_ANGLE_TOPIC, angleStr);
  } catch (...) {
    Serial.println("Exception in servo update");
  }
}

void sendAverageAndStatus() {
  try {
    if (sampleCount == 0) {
      Serial.println("No samples to average");
      return;
    }
    float sum = 0;
    for (int i = 0; i < sampleCount; i++) {
      sum += lightSamples[i];
    }
    float avgIntensity = sum / sampleCount;
    sampleCount = 0;
    char avgStr[8];
    snprintf(avgStr, sizeof(avgStr), "%.2f", avgIntensity);
    mqttClient.publish(LIGHT_AVG_TOPIC, avgStr);
    char statusMsg[120];
    int currentServoPos = servoInitialized ? servo.read() : -1;
    if (currentServoPos == -1) currentServoPos = (int)minAngle; // Fallback
    snprintf(statusMsg, sizeof(statusMsg), 
             "Light=%.2f, Temp=%.1f, Angle=%d, MinAngle=%d, Factor=%.2f, ServoInit=%s", 
             currentLightIntensity, currentTemperature, 
             currentServoPos, minAngle, controlFactor,
             servoInitialized ? "Y" : "N");
    mqttClient.publish("MEDIBOX/STATUS", statusMsg);
    Serial.print("Average light: ");
    Serial.println(avgIntensity);
  } catch (...) {
    Serial.println("Exception in sending average");
  }
}