#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <WiFi.h>
#include <PubSubClient.h>

// OLED
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Pins
#define TRIG 5
#define ECHO 18
#define FLOW_SENSOR 27
#define TURBIDITY_PIN 34
#define PH_PIN 35
#define TDS_PIN 32
#define VALVE 26

// Wi-Fi + MQTT
const char* ssid = "YOUR_WIFI";
const char* password = "YOUR_PASS";
const char* mqtt_server = "broker.hivemq.com";
WiFiClient espClient;
PubSubClient client(espClient);

// Flow measurement
volatile int flowPulseCount = 0;
float flowRate = 0;

void IRAM_ATTR flowPulse() {
  flowPulseCount++;
}

void reconnectMQTT() {
  while(!client.connected()){
    if(client.connect("WaterESP32")){
      client.subscribe("water/control");
    } else delay(2000);
  }
}

// Ultrasonic distance (water level)
long readDistanceCM() {
  digitalWrite(TRIG, LOW); delayMicroseconds(2);
  digitalWrite(TRIG, HIGH); delayMicroseconds(10);
  digitalWrite(TRIG, LOW);
  long duration = pulseIn(ECHO, HIGH);
  return duration * 0.034 / 2;
}

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

  // OLED
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.clearDisplay();
  display.setCursor(0,0);
  display.println("Water Distribution Init...");
  display.display();
  delay(2000);

  pinMode(TRIG, OUTPUT);
  pinMode(ECHO, INPUT);
  pinMode(FLOW_SENSOR, INPUT_PULLUP);
  pinMode(VALVE, OUTPUT);

  attachInterrupt(digitalPinToInterrupt(FLOW_SENSOR), flowPulse, RISING);

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

  client.setServer(mqtt_server,1883);
}

void loop() {
  if(!client.connected()) reconnectMQTT();
  client.loop();

  // --- Water Level ---
  long distance = readDistanceCM();
  float waterLevelPercent = map(distance,0,100,100,0);

  // --- Flow Rate ---
  flowRate = flowPulseCount * 2.25; // L/min approx
  flowPulseCount = 0;

  // --- Water Quality ---
  int turbidityRaw = analogRead(TURBIDITY_PIN); 
  int phRaw = analogRead(PH_PIN);               
  int tdsRaw = analogRead(TDS_PIN);

  float phValue = phRaw * 14.0 / 4095.0;                // pH scale 0-14
  float turbidityPercent = map(turbidityRaw,0,4095,0,100); // higher = dirtier
  float tdsValue = tdsRaw * 2000.0 / 4095.0;           // ppm (approx, module dependent)

  // --- Control Valve ---
  bool alert = false;
  if(waterLevelPercent < 20 || waterLevelPercent > 90){
    digitalWrite(VALVE, LOW); // stop pump
  } else {
    digitalWrite(VALVE, HIGH);
  }

  // --- Quality Alerts ---
  if(phValue < 6.5 || phValue > 8.5) alert = true;
  if(turbidityPercent > 70) alert = true;
  if(tdsValue > 1000) alert = true; // example threshold for TDS

  // --- OLED Display ---
  display.clearDisplay();
  display.setCursor(0,0);
  display.println("Water Distribution");
  display.print("Level %: "); display.println(waterLevelPercent);
  display.print("Flow L/min: "); display.println(flowRate);
  display.print("pH: "); display.println(phValue);
  display.print("Turbidity %: "); display.println(turbidityPercent);
  display.print("TDS ppm: "); display.println(tdsValue);
  display.print("Valve: "); display.println(digitalRead(VALVE)?"OPEN":"CLOSED");
  display.print("Alert: "); display.println(alert?"YES":"NO");
  display.display();

  // --- MQTT Publish ---
  String msg = "Level:"+String(waterLevelPercent)+
               " Flow:"+String(flowRate)+
               " pH:"+String(phValue)+
               " Turbidity:"+String(turbidityPercent)+
               " TDS:"+String(tdsValue)+
               " Valve:"+String(digitalRead(VALVE))+
               " Alert:"+String(alert);
  client.publish("water/status", msg.c_str());

  delay(1000);
}