# main.py – Plant monitoring system for Raspberry Pi Pico W

"""
OLED display docs: https://docs.micropython.org/en/latest/esp8266/tutorial/ssd1306.html
OLED fonts: https://github.com/easytarget/PrintPy2040
"""

import machine
from ssd1306 import SSD1306_I2C
from ezFBfont import ezFBfont
import ezFBfont_6x12_ascii_10 as thefont
import framebuf
import dht
import time
import plant_bitmaps as pb

# --- Constants ---
FRAME_COUNT = 27
RGB_RED_RANGE = (0, 8)
RGB_GREEN_RANGE = (9, 17)
RGB_BLUE_RANGE = (18, 26)
DEBOUNCE_MS = 50

# --- Pin assignments (match Wokwi wiring) ---
DHT_PIN = 2          # GP2  - DHT22 data
SOIL_ADC_PIN = 26    # GP26 - potentiometer used as "soil moisture"
LDR_ADC_PIN = 27     # GP27 - LDR AO
NTC_ADC_PIN = 28     # GP28 - NTC OUT

BTN_PIN = 3          # GP3  - pushbutton (WiFi config trigger, active-low)

LED_R_PIN = 14       # GP14 - RGB red
LED_G_PIN = 15       # GP15 - RGB green
LED_B_PIN = 16       # GP16 - RGB blue

LDC_SDA_PIN = 4      # GP4 - SDA (wired to oled1:SDA)
LDC_SCL_PIN = 5      # GP5 - SCL (wired to oled1:SCL)

# --- Display variables ---
PIX_RES_X = 128
PIX_RES_Y = 64
PIX_PADDING = 4

BITMAP_POS_X = PIX_RES_X - pb.WIDTH - PIX_PADDING  # right side image
BITMAP_POS_Y = PIX_PADDING

# --- Moisture calibration (DEFAULTS; override as needed) ---
SOIL_RAW_MIN = 0
SOIL_RAW_MAX = 65535
SOIL_FULLY_DRY_RAW = 15000
SOIL_FULLY_WET_RAW = 50000

# --- Peripherals ---
dht_sensor = dht.DHT22(machine.Pin(DHT_PIN))

soil_adc = machine.ADC(SOIL_ADC_PIN)
ldr_adc = machine.ADC(LDR_ADC_PIN)
ntc_adc = machine.ADC(NTC_ADC_PIN)

btn = machine.Pin(BTN_PIN, machine.Pin.IN, machine.Pin.PULL_UP)

led_r = machine.Pin(LED_R_PIN, machine.Pin.OUT)
led_g = machine.Pin(LED_G_PIN, machine.Pin.OUT)
led_b = machine.Pin(LED_B_PIN, machine.Pin.OUT)

# --- Button debounce state ---
btn_last_state = 1
btn_last_change = 0

# --- Helpers ---

def adc_to_volts(adc):
    """Convert ADC reading to voltage."""
    return adc.read_u16() * 3.3 / 65535

def clamp(value, lo, hi):
    """Clamp value between lo and hi."""
    if value < lo:
        return lo
    if value > hi:
        return hi
    return value

def map_linear(x, in_min, in_max, out_min, out_max):
    """Linear mapping from one range to another."""
    if in_max == in_min:
        return out_min
    return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

def set_soil_calibration(dry_raw=None, wet_raw=None, raw_min=None, raw_max=None):
    """
    Set soil moisture calibration values.
    Ensures dry_raw < wet_raw after setting.
    """
    global SOIL_FULLY_DRY_RAW, SOIL_FULLY_WET_RAW, SOIL_RAW_MIN, SOIL_RAW_MAX
    if raw_min is not None:
        SOIL_RAW_MIN = int(raw_min)
    if raw_max is not None:
        SOIL_RAW_MAX = int(raw_max)
    if dry_raw is not None:
        SOIL_FULLY_DRY_RAW = int(dry_raw)
    if wet_raw is not None:
        SOIL_FULLY_WET_RAW = int(wet_raw)
    if SOIL_FULLY_WET_RAW < SOIL_FULLY_DRY_RAW:
        SOIL_FULLY_DRY_RAW, SOIL_FULLY_WET_RAW = SOIL_FULLY_WET_RAW, SOIL_FULLY_DRY_RAW

def soil_raw_to_frame_index(soil_raw):
    """
    Convert raw soil moisture reading to animation frame index (0-26).
    Assumes SOIL_FULLY_DRY_RAW < SOIL_FULLY_WET_RAW (enforced by set_soil_calibration).
    """
    raw = clamp(soil_raw, SOIL_RAW_MIN, SOIL_RAW_MAX)
    raw = clamp(raw, SOIL_FULLY_DRY_RAW, SOIL_FULLY_WET_RAW)
    idx_f = map_linear(raw, SOIL_FULLY_DRY_RAW, SOIL_FULLY_WET_RAW, 0, FRAME_COUNT - 1)
    return max(0, min(FRAME_COUNT - 1, int(round(idx_f))))

def set_rgb(r=False, g=False, b=False):
    """
    Set RGB LED state.
    Common cathode RGB: 1 = on, 0 = off
    """
    led_r.value(1 if r else 0)
    led_g.value(1 if g else 0)
    led_b.value(1 if b else 0)

def set_rgb_by_frame(idx):
    """
    Set RGB LED color based on frame index.
    0-8 -> RED (dry), 9-17 -> GREEN (medium), 18-26 -> BLUE (wet)
    """
    if idx <= RGB_RED_RANGE[1]:
        set_rgb(r=True, g=False, b=False)
    elif idx <= RGB_GREEN_RANGE[1]:
        set_rgb(r=False, g=True, b=False)
    else:
        set_rgb(r=False, g=False, b=True)

def get_health_status(idx):
    """
    Get health status text based on frame index.
    0-8 -> DRY (red), 9-17 -> OK (green), 18-26 -> WET (blue)
    """
    if idx <= RGB_RED_RANGE[1]:
        return "DRY"
    elif idx <= RGB_GREEN_RANGE[1]:
        return "OK"
    else:
        return "WET"

def blit_bitmap(fb_dest, x, y, w, h, bitmap_data, color=1):
    """
    Blit a MONO_HLSB bitmap (horizontal, MSB first) into a MONO_VLSB framebuffer.
    
    The bitmap is stored row by row, with each row taking ceil(w/8) bytes.
    For 32x48: 4 bytes per row × 48 rows = 192 bytes total.
    Each byte represents 8 horizontal pixels, MSB = leftmost pixel.
    """
    bytes_per_row = (w + 7) // 8
    expected_size = bytes_per_row * h
    
    if len(bitmap_data) != expected_size:
        raise ValueError("Bitmap size mismatch: expected {} bytes, got {}".format(
            expected_size, len(bitmap_data)))
    
    # Process row by row
    for row in range(h):
        row_offset = row * bytes_per_row
        for col in range(w):
            byte_idx = row_offset + (col // 8)
            bit_idx = 7 - (col % 8)  # MSB first
            
            if (bitmap_data[byte_idx] >> bit_idx) & 1:
                fb_dest.pixel(x + col, y + row, color)

def draw_readings(fb, font, data, health_status):
    """
    Draw sensor readings to the framebuffer.
    """
    font.write("Health: {}".format(health_status), 0, 0)
    font.write("Soil: {}".format(data["soil_raw"]), 0, 14)
    font.write("Air: {:.1f}C".format(data["air_temp"]), 0, 24)
    font.write("Hum: {:.1f}%".format(data["humidity"]), 0, 34)
    font.write("Light: {}".format(data["ldr_raw"]), 0, 44)
    font.write("Temp: {}".format(data["air_temp"]), 0, 54)

def validate_sensor_data(data):
    """
    Validate sensor readings to detect errors.
    DHT22 can return invalid readings.
    Returns True if data is valid, False otherwise.
    """
    # Check for DHT22 error values
    if data["air_temp"] < -40 or data["air_temp"] > 80:
        return False
    if data["humidity"] < 0 or data["humidity"] > 100:
        return False
    return True

def read_sensors():
    """
    Read all sensors and return a dictionary of values.
    Includes basic error handling for DHT22.
    """
    try:
        dht_sensor.measure()
        air_temp = dht_sensor.temperature()
        humidity = dht_sensor.humidity()
    except OSError as e:
        print("DHT22 read error:", e)
        air_temp = 0.0
        humidity = 0.0

    soil_raw = soil_adc.read_u16()
    ldr_raw = ldr_adc.read_u16()
    ntc_raw = ntc_adc.read_u16()
    ntc_voltage = adc_to_volts(ntc_adc)

    return {
        "air_temp": air_temp,
        "humidity": humidity,
        "soil_raw": soil_raw,
        "ldr_raw": ldr_raw,
        "ntc_raw": ntc_raw,
        "ntc_voltage": ntc_voltage,
    }

def check_button():
    """
    Check button state with debouncing.
    Returns True if button was pressed (with debounce), False otherwise.
    """
    global btn_last_state, btn_last_change
    
    current_time = time.ticks_ms()
    btn_state = btn.value()
    
    # Detect falling edge (button press) with debounce
    if btn_state == 0 and btn_last_state == 1:
        if time.ticks_diff(current_time, btn_last_change) > DEBOUNCE_MS:
            btn_last_change = current_time
            btn_last_state = btn_state
            return True
    
    btn_last_state = btn_state
    return False

# --- Main loop ---

def main():
    """Main application loop."""
    # I2C0 on GP4/GP5
    i2c = machine.I2C(0, sda=machine.Pin(LDC_SDA_PIN), scl=machine.Pin(LDC_SCL_PIN), freq=400000)
    print("I2C devices found:", [hex(a) for a in i2c.scan()])  # should include 0x3c
    oled = SSD1306_I2C(128, 64, i2c)     # addr defaults to 0x3c

    # Off-screen double buffer (MONO_VLSB)
    buf = bytearray(PIX_RES_X * PIX_RES_Y // 8)  # 1024 bytes
    fb = framebuf.FrameBuffer(buf, PIX_RES_X, PIX_RES_Y, framebuf.MONO_VLSB)

    # Font that targets the off-screen buffer
    font = ezFBfont(fb, thefont, vgap=0, verbose=True)

    # Initial splash screen
    print("Initializing display...")
    fb.fill(0)
    
    # Draw start screen centered
    font.write("Plant Buddy", PIX_RES_X // 2, 20, halign='center')
    font.write("v1.0.0", PIX_RES_X // 2, 32, halign='center')
    
    # Show splash for 2 seconds
    oled.blit(fb, 0, 0)
    oled.show()
    print("Display initialized. Starting main loop...")
    time.sleep(2)

    # Main sensor reading loop
    while True:
        try:
            data = read_sensors()
            
            # Validate sensor data
            if not validate_sensor_data(data):
                print("Warning: Invalid sensor data detected")

            # Prepare next frame in off-screen buffer - clear everything
            fb.fill(0)

            # Plant bitmap based on current soil_raw
            idx = soil_raw_to_frame_index(data["soil_raw"])
            frame = pb.plant_images[idx]
            blit_bitmap(fb, BITMAP_POS_X, BITMAP_POS_Y, pb.WIDTH, pb.HEIGHT, frame, color=1)

            # Get health status based on moisture level
            health_status = get_health_status(idx)

            # Draw sensor readings (no static header in main loop)
            draw_readings(fb, font, data, health_status)

            # Blit once per loop (no flicker)
            oled.blit(fb, 0, 0)
            oled.show()

            # RGB LED reflects same moisture level
            set_rgb_by_frame(idx)

            # Console log
            print(
                "Air: {:.1f}C  Hum: {:.1f}%  Soil: {}  Light: {}  NTC: {} ({:.2f}V)  Frame: {}".format(
                    data["air_temp"],
                    data["humidity"],
                    data["soil_raw"],
                    data["ldr_raw"],
                    data["ntc_raw"],
                    data["ntc_voltage"],
                    idx,
                )
            )

            # Button for WiFi config mode (with debouncing)
            if check_button():
                print("Button pressed -> enter WiFi config mode here")
                # Add WiFi config mode logic here

            time.sleep(0.5)

        except Exception as e:
            print("Error in main loop:", e)
            import sys
            sys.print_exception(e)
            time.sleep(1)


if __name__ == '__main__':
    main()