from machine import Pin, I2C, PWM
from neopixel import NeoPixel
import ssd1306
import dht
import time
import network
import BlynkLib
import gc
import machine

# ================== INITIAL SETUP ==================
gc.enable()

# ================== BLYNK & WIFI SETUP ==================
WIFI_SSID = "MiMikyu" 
WIFI_PASS = "pattarin192"
BLYNK_AUTH = "0f8LcG45QzPkWYQXVAae7pdUvN4J0UcL"

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        wlan.connect(WIFI_SSID, WIFI_PASS)
        for _ in range(20): 
            if wlan.isconnected(): break
            time.sleep(0.5)
    gc.collect()

connect_wifi()
blynk = BlynkLib.Blynk(BLYNK_AUTH, insecure=True)

# ================== STABILITY GUARDS ==================
last_wifi_check = 0
last_blynk_check = 0
last_watchdog_kick = time.ticks_ms()
WATCHDOG_TIMEOUT = 30000 
last_sent_color = None # 3️⃣ กัน Blynk ส่งสีซ้ำ

# ================== HARD WDT (SAFE ADD) ==================
try:
    wdt = machine.WDT(timeout=WATCHDOG_TIMEOUT)
except:
    wdt = None

def wifi_reconnect_guard():
    wlan = network.WLAN(network.STA_IF)
    if not wlan.isconnected():
        try: wlan.connect(WIFI_SSID, WIFI_PASS)
        except: pass

def blynk_reconnect_guard():
    try:
        if not blynk.connected():
            blynk.connect()
    except: pass

def watchdog_check():
    if time.ticks_diff(time.ticks_ms(), last_watchdog_kick) > WATCHDOG_TIMEOUT:
        machine.reset()

# ================== DEVICE CONFIG ==================
NUM_PIXELS = 16
PIXEL_PIN = 25
NP_BTN_PIN = 32
np = NeoPixel(Pin(PIXEL_PIN, Pin.OUT), NUM_PIXELS)
np_btn = Pin(NP_BTN_PIN, Pin.IN, Pin.PULL_UP)

fixed_colors = [
    (255, 0, 128), (0, 255, 255), (0, 64, 128), (135, 206, 250),
    (150, 150, 255), (0, 128, 64), (152, 251, 152), (0, 128, 128),
    (128, 64, 0), (255, 218, 185), (230, 230, 250), (255, 255, 224),
    "relax", "rainbow_static", "rainbow_scroll", "auto_cycle"
]
color_names = ["Pink", "Cyan", "Blue Cool", "Light Blue", "Blue Purple", "Mint Green", "Light Green", "Teal", "Orange Warm", "Peach", "Lavender", "Light Yellow", "Relax", "Rainbow", "Scroll", "Auto"]
relax_colors = [(0, 64, 128), (135, 206, 250), (150, 150, 255), (0, 128, 64), (152, 251, 152), (0, 128, 128)]

# --- Variables ---
color_index = 0
brightness_step = 4 
brightness_levels = [0, 64, 128, 192, 255]
np_on = True
np_hold_start = None
last_np_btn = 1
click_time = 0
click_count = 0
rainbow_pos = 0
relax_pos = 0
auto_pos = 0
RELAX_SPEED = 500
AUTO_SPEED = 1000
last_relax_time = 0
last_auto_time = 0
device_on = True 
page = 0
custom_rgb = None
blink_state = False 
last_v12_update = 0
V12_UPDATE_INTERVAL = 150 
temp = hum = 0 

# ================== OLED & SENSORS ==================
i2c = I2C(0, scl=Pin(22), sda=Pin(21))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)
led = Pin(4, Pin.OUT)
led.value(1)
button = Pin(13, Pin.IN, Pin.PULL_UP)
power_button = Pin(15, Pin.IN, Pin.PULL_UP)
buzzer = PWM(Pin(27))
buzzer.duty_u16(0)
sensor = dht.DHT22(Pin(14))
relay = Pin(26, Pin.OUT) 
relay_master = True
relay_mode = 1           
relay_interval = 3000
last_relay_time = 0
relay_state = False

# ================== HELPER FUNCTIONS ==================
def wheel(pos):
    if pos < 85: return (255 - pos*3, pos*3, 0)
    elif pos < 170: pos -= 85; return (0, 255 - pos*3, pos*3)
    else: pos -= 170; return (pos*3, 0, 255 - pos*3)

def apply_brightness(color, level):
    val = brightness_levels[level]
    factor = val / 255
    return (int(color[0]*factor), int(color[1]*factor), int(color[2]*factor))

def beep(duration=0.03, freq=1200):
    try:
        buzzer.freq(freq); buzzer.duty_u16(20000); time.sleep(duration); buzzer.duty_u16(0)
    except: pass

# ================== DRAW PAGES ==================
def refresh_display():
    if page == 0: draw_page_main()
    elif page == 1: draw_page_dht(temp, hum)
    elif page == 2: draw_page_np()
    elif page == 3: draw_page_relay()

def draw_page_main():
    oled.fill(0); oled.text("Smart", 44, 5); oled.text("Therapy Box", 20, 18)
    oled.text("By Miss PATTARIN", 0, 35); oled.text("LED: ON", 0, 50); oled.show()

def draw_page_dht(t, h, show_overheat=True):
    oled.fill(0)
    if t >= 30 and show_overheat:
        oled.fill_rect(0, 0, 128, 12, 1); oled.text("OVERHEAT!", 30, 2, 0)
    else: oled.text("DHT SENSOR", 25, 0)
    oled.text("Temp: {:.1f} C".format(t), 0, 25)
    oled.text("Hum : {:.1f} %".format(h), 0, 45); oled.show()

def draw_page_np():
    oled.fill(0); oled.text("NEOPIXEL CTRL", 0, 0)
    name = color_names[color_index] if custom_rgb is None else "Custom"
    oled.text("Color: {}".format(name), 0, 20)
    oled.text("Bright: {}%".format(brightness_step*25), 0, 40)
    oled.text("Status: {}".format('ON' if np_on else 'OFF'), 0, 55); oled.show()

def draw_page_relay():
    oled.fill(0); oled.text("RELAY CONTROL", 8, 0)
    oled.text("Master: {}".format('ON' if relay_master else 'OFF'), 0, 20)
    oled.text("Mode: {}".format(['ON','INTERVAL'][relay_mode]), 0, 35)
    oled.text("Every: {}s".format(relay_interval // 1000), 0, 50); oled.show()

# ================== CORE LOGIC ==================
def send_color_feedback(rgb_color):
    global last_v12_update, last_sent_color
    now = time.ticks_ms()
    try:
        final_color = apply_brightness(rgb_color, brightness_step)
        hex_color = "#%02x%02x%02x" % final_color
        
        # 3️⃣ กัน Blynk ส่งสีซ้ำ (Guard)
        if last_sent_color == hex_color:
            return
        last_sent_color = hex_color
        
        blynk.set_property(12, "color", hex_color)
        blynk.virtual_write(12, 255 if np_on else 0)
        last_v12_update = now
    except: pass

def update_neopixel(force_blynk=False):
    global rainbow_pos, relax_pos, auto_pos, last_relax_time, last_auto_time, last_v12_update
    now = time.ticks_ms()
    
    if not np_on or not device_on:
        np.fill((0,0,0)); np.write()
        if force_blynk or time.ticks_diff(now, last_v12_update) > V12_UPDATE_INTERVAL:
            try: blynk.virtual_write(12, 0)
            except: pass
            last_v12_update = now
        return

    current_rgb = (0,0,0)
    feedback_rgb = (0,0,0)
    is_animated = False

    if custom_rgb:
        current_rgb = custom_rgb; feedback_rgb = custom_rgb
    else:
        mode = fixed_colors[color_index]
        if mode == "relax":
            if time.ticks_diff(now, last_relax_time) > RELAX_SPEED:
                relax_pos = (relax_pos + 1) % len(relax_colors); last_relax_time = now
            feedback_rgb = relax_colors[relax_pos]; current_rgb = feedback_rgb
            is_animated = True
        elif mode == "rainbow_static":
            for i in range(NUM_PIXELS): np[i] = apply_brightness(wheel(i*256//NUM_PIXELS), brightness_step)
            feedback_rgb = wheel(rainbow_pos) 
        elif mode == "rainbow_scroll":
            for i in range(NUM_PIXELS): np[i] = apply_brightness(wheel((i*8 + rainbow_pos) & 255), brightness_step)
            feedback_rgb = wheel(rainbow_pos); rainbow_pos = (rainbow_pos + 5) & 255
            is_animated = True
        elif mode == "auto_cycle":
            if time.ticks_diff(now, last_auto_time) > AUTO_SPEED:
                auto_pos = (auto_pos + 1) % len(relax_colors); last_auto_time = now
            feedback_rgb = relax_colors[auto_pos]; current_rgb = feedback_rgb
            is_animated = True
        else:
            current_rgb = mode; feedback_rgb = mode

    if not (fixed_colors[color_index] in ["rainbow_static", "rainbow_scroll"]):
        np.fill(apply_brightness(current_rgb, brightness_step))
    np.write()

    if force_blynk or (is_animated and time.ticks_diff(now, last_v12_update) > V12_UPDATE_INTERVAL):
        send_color_feedback(feedback_rgb)

def start_all_devices():
    global np_on, relay_master, device_on, page
    beep(); led.value(1); np_on = True; device_on = True; page = 0
    update_neopixel(force_blynk=True); relay_master = True; relay.value(1)
    try: sensor.measure()
    except: pass
    draw_page_main()
    try:
        blynk.set_property(11, "color", "#00FF00")
        blynk.virtual_write(11, 255)
        blynk.virtual_write(1, 1); blynk.virtual_write(0, 0)
        blynk.virtual_write(2, brightness_step); blynk.virtual_write(3, color_index)
        blynk.virtual_write(9, 1); blynk.virtual_write(6, 1); blynk.virtual_write(10, relay_mode)
        blynk.virtual_write(15, 0)
        blynk.virtual_write(8, relay_interval // 1000)
    except: pass

# ================== BLYNK HANDLERS ==================
@blynk.on("connected")
def blynk_connected():
    try:
        if not device_on:
            blynk.set_property(11, "color", "#FF0000")
            blynk.virtual_write(11, 255)
        else:
            blynk.set_property(11, "color", "#00FF00")
            blynk.virtual_write(11, 255)
    except: pass
    blynk.sync_virtual(1,2,3,6,8,9,10,15)

@blynk.on("V1")
def blynk_power(v):
    global device_on, page
    if int(v[0]): 
        page = 0; start_all_devices()
    else: 
        device_on = False; relay.value(0); np.fill((0,0,0)); np.write()
        oled.fill(0); oled.text("Power OFF", 25, 28); oled.show()
        try: 
            blynk.set_property(11, "color", "#FF0000")
            blynk.virtual_write(11, 255)
            blynk.virtual_write(0, 0); blynk.virtual_write(12, 0)
        except: pass

@blynk.on("V2")
def blynk_bright(v):
    global brightness_step, page
    if not device_on: return
    brightness_step = int(v[0]); update_neopixel(force_blynk=True)
    page = 2; blynk.virtual_write(15, 2); draw_page_np()

@blynk.on("V3")
def blynk_menu(v):
    global color_index, custom_rgb, page
    if not device_on: return
    color_index = int(v[0]); custom_rgb = None; update_neopixel(force_blynk=True)
    page = 2; blynk.virtual_write(15, 2); draw_page_np()

@blynk.on("V6")
def blynk_relay_master(v):
    global relay_master, page
    if not device_on: return
    relay_master = bool(int(v[0]))
    if not relay_master: relay.value(0)
    page = 3; blynk.virtual_write(15, 3); draw_page_relay()

@blynk.on("V7")
def blynk_zergba(v):
    global custom_rgb, page
    if not device_on: return
    try:
        custom_rgb = (int(v[0]), int(v[1]), int(v[2]))
        update_neopixel(force_blynk=True)
        page = 2; blynk.virtual_write(15, 2); draw_page_np()
    except: pass

@blynk.on("V8")
def blynk_interval(v):
    global relay_interval, page
    if not device_on: return
    val = int(v[0])
    relay_interval = max(1, min(20, val)) * 1000
    page = 3; blynk.virtual_write(15, 3); draw_page_relay()

@blynk.on("V9")
def blynk_np_toggle(v):
    global np_on, page
    if not device_on: return
    np_on = bool(int(v[0])); update_neopixel(force_blynk=True)
    page = 2; blynk.virtual_write(15, 2); draw_page_np()

@blynk.on("V10")
def blynk_relay_mode(v):
    global relay_mode, page
    if not device_on: return
    relay_mode = int(v[0])
    page = 3; blynk.virtual_write(15, 3); draw_page_relay()

@blynk.on("V15")
def blynk_change_page(v):
    global page
    if not device_on: return
    new_page = int(v[0])
    if new_page != page:
        page = new_page
        beep(0.02); refresh_display()

# ================== MAIN LOOP ==================
last_state = last_power_state = last_np_btn = 1
last_dht_read = last_blink_time = 0

# เริ่มต้น: ตั้งค่า V11 เป็นสีขาว (#FFFFFF)
try:
    blynk.set_property(11, "color", "#FFFFFF")
    blynk.virtual_write(11, 255)
except: pass

start_all_devices()

while True:
    now = time.ticks_ms()
    
    # 2️⃣ คิก WDT ใน loop (Hard Feed)
    if wdt:
        try: wdt.feed()
        except: pass
    
    try: blynk.run()
    except: pass
    
    if time.ticks_diff(now, last_wifi_check) > 10000:
        last_wifi_check = now; wifi_reconnect_guard()
    if time.ticks_diff(now, last_blynk_check) > 15000:
        last_blynk_check = now; blynk_reconnect_guard()
    last_watchdog_kick = now; watchdog_check()

    p = power_button.value()
    if last_power_state == 1 and p == 0:
        device_on = not device_on; beep()
        if device_on: 
            start_all_devices()
        else:
            oled.fill(0); oled.text("Power OFF", 25, 28); oled.show()
            np.fill((0,0,0)); np.write(); relay.value(0)
            try: 
                blynk.set_property(11, "color", "#FF0000")
                blynk.virtual_write(11, 255)
                blynk.virtual_write(1, 0); blynk.virtual_write(12, 0)
            except: pass
    last_power_state = p

    if not device_on:
        time.sleep(0.01); continue

    update_neopixel()

    np_val = np_btn.value()
    if np_val == 0 and last_np_btn == 1: np_hold_start = now
    elif np_val == 0 and np_hold_start is not None:
        if time.ticks_diff(now, np_hold_start) > 1000:
            np_on = not np_on; update_neopixel(force_blynk=True)
            try: blynk.virtual_write(9, 1 if np_on else 0)
            except: pass
            page = 2; blynk.virtual_write(15, 2); draw_page_np()
            np_hold_start = None; click_count = 0
    elif np_val == 1 and last_np_btn == 0:
        if np_hold_start is not None and time.ticks_diff(now, np_hold_start) < 1000:
            click_count += 1; click_time = now
        np_hold_start = None
    last_np_btn = np_val

    if click_count > 0 and time.ticks_diff(now, click_time) > 250:
        if page == 3:
            if click_count == 1: 
                relay_mode = (relay_mode + 1) % 2
                try: blynk.virtual_write(10, relay_mode)
                except: pass
            elif click_count == 2:
                relay_master = not relay_master
                try: blynk.virtual_write(6, 1 if relay_master else 0)
                except: pass
            draw_page_relay()
        elif page == 2:
            if click_count == 1: 
                color_index = (color_index + 1) % len(fixed_colors); custom_rgb = None
                try: blynk.virtual_write(3, color_index)
                except: pass
            elif click_count == 2: 
                brightness_step = (brightness_step + 1) % 5
                try: blynk.virtual_write(2, brightness_step)
                except: pass
            update_neopixel(force_blynk=True); draw_page_np()
        click_count = 0

    s = button.value()
    if last_state == 1 and s == 0:
        beep(); page = (page + 1) % 4
        try: blynk.virtual_write(15, page)
        except: pass
        refresh_display()
    last_state = s

    if time.ticks_diff(now, last_dht_read) > 2000:
        last_dht_read = now
        try:
            sensor.measure(); temp = sensor.temperature(); hum = sensor.humidity()
            try:
                blynk.virtual_write(4, "{:.1f}".format(temp))
                blynk.virtual_write(5, "{:.1f}".format(hum))
            except: pass
            if page == 1: draw_page_dht(temp, hum)
        # 4️⃣ กัน DHT error ยาว ๆ (Reset 0)
        except:
            temp = hum = 0

    if temp >= 30:
        if page != 1: 
            page = 1
            try: blynk.virtual_write(15, 1)
            except: pass
        if time.ticks_diff(now, last_blink_time) > 300:
            blink_state = not blink_state; last_blink_time = now
            try: blynk.virtual_write(0, 255 if blink_state else 0)
            except: pass
            led.value(blink_state); draw_page_dht(temp, hum, show_overheat=blink_state); beep(0.05)
    else: 
        led.value(1); 
        try: blynk.virtual_write(0, 0)
        except: pass

    if relay_master:
        if relay_mode == 0: relay.value(1)
        else:
            if time.ticks_diff(now, last_relay_time) > relay_interval:
                relay_state = not relay_state; last_relay_time = now
            relay.value(relay_state)
    else: relay.value(0)
    
    # 5️⃣ ลด GC สุ่มค้าง (Scheduled Collect ทุก ~5 วินาที)
    if now % 5000 < 20:
        gc.collect()

    time.sleep(0.01)