import time
import cv2
import numpy as np
import threading
from ultralytics import YOLO
import paho.mqtt.client as mqtt
import os
from gpiozero import DistanceSensor
from luma.core.interface.serial import i2c
from luma.oled.device import ssd1306
from PIL import Image, ImageDraw, ImageFont

# =================== CONFIGURATION ===================
MODEL_PATH = "modelpath"
RESOLUTION_W = 1280
RESOLUTION_H = 720
CONFIDENCE_THRESHOLD = 0.6
MQTT_BROKER = "127.0.0.1" 
MQTT_PORT = 1883
COMMAND_FILE = "path to command file"
MQTT_TOPICS = ["local", "lamp1/broadcast", "lamp2/broadcast"]
TARGET_NAMES = ["Main", "Lamp 1", "Lamp 2"]
current_target = 0

# Sensor parameters
DISTANCE_MIN = 0.4  # 40 cm
DISTANCE_MAX = 0.7  # 70 cm
ACTIVATION_THRESHOLD = 0.15  # 15 cm change
SUSTAIN_TIME = 2.0  # 2 seconds

# Gesture timing
DEBOUNCE_INTERVAL = 2.0
FIST_HOLD_TIME = 2.0
COOLDOWN_AFTER_FIST = 1.0
GESTURE_STABILITY_TIME = 0.3  # 300ms stability check

# OLED update control
DISPLAY_UPDATE_INTERVAL = 0.2  # 5 times per second

# Bounding box colors
colors = [
    (0, 255, 0),     # Green
    (0, 0, 255),     # Red
    (255, 0, 0),     # Blue
    (255, 255, 0),   # Cyan
    (0, 255, 255),   # Yellow
    (255, 0, 255),   # Magenta
    (128, 0, 0),     # Navy
    (0, 128, 0),     # Dark Green
    (128, 128, 0),   # Olive
    (0, 0, 128)      # Maroon
]

# Set environment variable for X11 display
os.environ["DISPLAY"] = ":0"

# =====================================================

# MQTT Setup
def on_publish(client, userdata, mid):
    print("[MQTT] Message published")

def publish_to_target(msg):
    topic = MQTT_TOPICS[current_target]
    mqtt_client.publish(topic, msg.encode('utf-8'))
    print(f"[MQTT] Sent '{msg}' to {topic}")

mqtt_client = mqtt.Client(client_id="rpi_ai", protocol=mqtt.MQTTv5)
mqtt_client.on_publish = on_publish
mqtt_client.connect(MQTT_BROKER, MQTT_PORT)
mqtt_client.loop_start()

# YOLO Setup
print(f"Loading model from {MODEL_PATH}...")
model = YOLO(MODEL_PATH, task='detect')
labels = model.names

# Camera Setup
cap = cv2.VideoCapture(0)
cap.set(3, RESOLUTION_W)
cap.set(4, RESOLUTION_H)
if not cap.isOpened():
    print("Error: Cannot open camera")
    exit()

# OLED Setup
serial = i2c(port=2, address=0x3C)
device = ssd1306(serial, width=128, height=64)
font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 14)
font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 12)

# Distance Sensor Setup
sensor = DistanceSensor(echo=24, trigger=23)
activation_checking = False

# Gesture stability variables
last_gesture = None
last_gesture_time = 0

# Display state tracking
last_display = {
    "distance": None,
    "gesture": None,
    "lamp_on": None,
    "target": None
}
last_update_time = 0

def display_minimal(distance, gesture, lamp_on, target):
    global last_display, last_update_time
    current_time = time.time()
    
    if (last_display["distance"] == distance and
        last_display["gesture"] == gesture and
        last_display["lamp_on"] == lamp_on and
        last_display["target"] == target and
        (current_time - last_update_time) < DISPLAY_UPDATE_INTERVAL):
        return

    last_display.update({
        "distance": distance,
        "gesture": gesture,
        "lamp_on": lamp_on,
        "target": target
    })
    last_update_time = current_time

    image = Image.new("1", (device.width, device.height))
    draw = ImageDraw.Draw(image)
    # First row: Distance and Target
    draw.text((0, 0), f"{distance*100:.0f}cm", font=font_large, fill=255)
    draw.text((70, 0), target, font=font_small, fill=255)
    # Second row: Gesture status
    gesture_text = gesture if gesture else "No gesture"
    draw.text((0, 24), gesture_text, font=font_large, fill=255)
    # Third row: Lamp status
    draw.text((0, 48), f"Lamp: {'ON' if lamp_on else 'OFF'}", font=font_large, fill=255)
    device.display(image)

def get_stable_gesture(detected_gesture, current_time):
    global last_gesture, last_gesture_time
    if detected_gesture == last_gesture:
        last_gesture_time = current_time
        return detected_gesture
    else:
        if last_gesture is None:
            last_gesture = detected_gesture
            last_gesture_time = current_time
            return None
        else:
            if current_time - last_gesture_time >= GESTURE_STABILITY_TIME:
                last_gesture = detected_gesture
                last_gesture_time = current_time
                return detected_gesture
            else:
                return None

def check_abrupt_activation():
    global lamp_on
    while activation_checking:
        if not lamp_on:
            initial_distance = sensor.distance
            start_time = None
            activated = False
            
            while not activated and activation_checking and not lamp_on:
                current_distance = sensor.distance
                if abs(current_distance - initial_distance) > ACTIVATION_THRESHOLD:
                    start_time = start_time or time.time()
                    if time.time() - start_time >= SUSTAIN_TIME:
                        publish_to_target("palm")
                        lamp_on = True
                        activated = True
                else:
                    start_time = None
                time.sleep(0.1)
        time.sleep(0.5)

# Main Loop Vars
frame_count = 0
start_time = time.time()
last_command = ""
last_sent_gesture = None
last_sent_time = 0
gesture_cooldown_until = 0
fist_held_start = None
lamp_on = False

print("\nStarting detection...")
print("Press Ctrl+C to stop\n")
print("=" * 50)

try:
    activation_checking = True
    activation_thread = threading.Thread(target=check_abrupt_activation)
    activation_thread.start()

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Error: Camera read failed")
            break

        frame_count += 1
        current_time = time.time()

        # Distance processing
        current_distance = sensor.distance
        in_range = DISTANCE_MIN <= current_distance <= DISTANCE_MAX

        # Gesture detection (only when in range)
        current_gesture = None
        detections = []
        if in_range:
            results = model(frame, verbose=False)
            detections = results[0].boxes
            for det in detections:
                if det.conf.item() > CONFIDENCE_THRESHOLD:
                    current_gesture = labels[int(det.cls.item())]
                    break

        # Apply gesture stability check
        current_gesture = get_stable_gesture(current_gesture, current_time)

        if current_time < gesture_cooldown_until:
            current_gesture = None

        # Draw bounding boxes
        if len(detections) > 0:
            for det in detections:
                classidx = int(det.cls.item())
                conf = det.conf.item()
                if conf > CONFIDENCE_THRESHOLD:
                    box = det.xyxy[0].cpu().numpy()
                    x1, y1, x2, y2 = map(int, box)
                    color = colors[classidx % len(colors)]
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    label = f"{labels[classidx]} {conf*100:.1f}%"
                    (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                    cv2.rectangle(frame, (x1, y1-20), (x1+label_width, y1), color, -1)
                    cv2.putText(frame, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        # Display FPS on frame
        fps_text = f"FPS: {1/(current_time - start_time):.1f}" if (current_time - start_time) > 0 else "FPS: N/A"
        cv2.putText(frame, fps_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

        # Show the camera feed
        cv2.imshow('Smart Lamp Camera Feed', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        # Fist hold handling
        if current_gesture == "fist":
            fist_held_start = fist_held_start or current_time
            if current_time - fist_held_start >= FIST_HOLD_TIME:
                current_target = (current_target + 1) % len(MQTT_TOPICS)
                print(f"Switched target to: {TARGET_NAMES[current_target]}")
                gesture_cooldown_until = current_time + COOLDOWN_AFTER_FIST
                fist_held_start = None
        else:
            fist_held_start = None

        # Gesture processing
        if current_gesture in {"palm", "like", "dislike", "vsign"}:
            if current_gesture != last_sent_gesture or (current_time - last_sent_time > DEBOUNCE_INTERVAL):
                publish_to_target(current_gesture)
                last_sent_gesture = current_gesture
                last_sent_time = current_time
                gesture_cooldown_until = current_time + DEBOUNCE_INTERVAL

        # Display update
        display_minimal(current_distance, current_gesture, lamp_on, TARGET_NAMES[current_target])

        # External commands
        if os.path.exists(COMMAND_FILE):
            with open(COMMAND_FILE, "r") as f:
                command = f.read().strip()
            if command and command != last_command:
                mqtt_client.publish(MQTT_TOPICS[current_target], command.encode())
                last_command = command
                lamp_on = (command == "palm")

        # Performance monitoring
        if (current_time - start_time) >= 10:
            print(f"FPS: {frame_count / (current_time - start_time):.1f}")
            frame_count = 0
            start_time = current_time

        time.sleep(0.05)

except KeyboardInterrupt:
    print("\nStopping...")

finally:
    activation_checking = False
    activation_thread.join()
    cap.release()
    mqtt_client.loop_stop()
    mqtt_client.disconnect()
    device.clear()
    cv2.destroyAllWindows()
    print("Cleanup complete. Goodbye!")