"""
automation.py — Crate-Scan Pro autonomous conveyor state machine.

Runs as conveyor.service, independent of the Gradio UI (cratescan.service):

    WAITING --(crate <= 26cm)--> RUNNING_TO_CAMERA --(crate <= 15cm)--> SCANNING --> EJECTING --> WAITING

GPIO 5 (push button) toggles WAITING/RUNNING_TO_CAMERA/SCANNING/EJECTING <-> IDLE.

This process is 100% autonomous: it owns the conveyor motor, the distance
sensor, the RGB LED, and the 12V LED strips directly via gpiozero. It makes
exactly two Gradio API calls — capture_usb_camera_frame and run_yolo — and
only because the camera + YOLO model live in app.py's process; those are
data/compute calls, not physical hardware control. If app.py is offline or
slow, those calls fail gracefully (see _gradio_connect_loop) and the
conveyor keeps running — no actuator in this file ever depends on app.py.

PIN CONFIGURATION  —  BCM numbers, matches RPi/debugging.py wiring
"""

import time
import threading
from collections import deque

try:
    from gradio_client import Client, handle_file
    GRADIO_CLIENT_AVAILABLE = True
except ImportError:
    GRADIO_CLIENT_AVAILABLE = False

from gpiozero.pins.lgpio import LGPIOFactory
from gpiozero import Device, DistanceSensor, Motor, RGBLED, Button

Device.pin_factory = LGPIOFactory()

# ── PIN CONFIGURATION  (matches RPi/debugging.py) ────────────────────────────
TRIGGER_PIN   = 23   # HC-SR04 distance sensor
ECHO_PIN      = 24

CONVEYOR_IN1  = 8     # L298N channel driving the conveyor belt motor
CONVEYOR_IN2  = 7

LED_STRIP_IN1 = 10    # L298N channel driving the 2x 12V LED strips
LED_STRIP_IN2 = 9

RGB_RED_PIN   = 17
RGB_GREEN_PIN = 27
RGB_BLUE_PIN  = 22

PAUSE_BUTTON_PIN = 5

RGB_COLORS = {
    "red":    (1, 0, 0),
    "green":  (0, 1, 0),
    "yellow": (1, 1, 0),
    "blue":   (0, 0, 1),
    "purple": (0.5, 0, 0.5),   # IDLE indicator
    "white":  (1, 1, 1),       # active / running indicator
}

# ── DISTANCE THRESHOLDS (cm) ──────────────────────────────────────────────────
ENTRY_DISTANCE_CM  = 26   # crate has entered the tunnel  -> start conveyor
CAMERA_DISTANCE_CM = 15   # crate is directly under the camera -> stop conveyor

# ── TIMING ─────────────────────────────────────────────────────────────────
LED_WARMUP_SECONDS     = 2.5   # minimum dwell time under the lights before the shot
RESULT_COLOR_SECONDS   = 10.0  # how long the red/yellow/green result shows before reverting
EJECT_TIMEOUT_SECONDS  = 8.0   # safety cap in case the crate never clears the sensor
POLL_INTERVAL_SECONDS  = 0.1
DISTANCE_HISTORY_SIZE  = 5     # moving-average window (in poll cycles) over the distance reading
HARDWARE_RETRY_SECONDS = 2.0   # how often the healing loop retries a busy GPIO device
GRADIO_CONNECT_TIMEOUT = 5     # seconds — bounds the handshake so a slow/offline app.py can't hang this process
GRADIO_RETRY_SECONDS   = 5.0   # how often the background loop retries connecting to the Gradio API
PAUSE_DEBOUNCE_SECONDS = 0.5   # ignore bounce-induced double-fires so one physical press = one toggle

# capture_usb_camera_frame and run_yolo are two separate Gradio events (the
# manual flow is Capture Frame -> Run YOLO) — data/compute calls only, since
# the camera + YOLO model live in app.py's process. See module docstring.
GRADIO_URL = "http://127.0.0.1:7860"

# ── STATE MACHINE STATES ─────────────────────────────────────────────────────
WAITING           = "WAITING"
RUNNING_TO_CAMERA = "RUNNING_TO_CAMERA"
SCANNING          = "SCANNING"
EJECTING          = "EJECTING"
IDLE              = "IDLE"   # entered via the pause button — RGB LED goes purple


class CrateScanAutomation:
    """Owns every GPIO device directly and drives the state machine."""

    def __init__(self):
        self.conveyor = None
        self.distance_sensor = None
        self.rgb_led = None
        self.led_strips = None
        self._distance_history = deque(maxlen=DISTANCE_HISTORY_SIZE)

        # Create button WITHOUT the callback registered yet.
        # Registering when_pressed early can cause a spurious IDLE fire before
        # force_initial_state() has a chance to read the real pin level.
        self.pause_button = Button(PAUSE_BUTTON_PIN, bounce_time=0.1)

        self.state = WAITING
        self._previous_state = None    # state to restore when we leave IDLE
        self._eject_started_at = None
        self._last_toggle_time = 0.0   # guards against bounce-induced double-fires
        self.last_button_state = False  # rising-edge tracker for toggle behaviour
        self._lock = threading.Lock()  # guards self.state / self.conveyor / self.distance_sensor
        self._running = True

        # The Gradio client (used only for capture_usb_camera_frame + run_yolo,
        # never for actuators) connects on its own thread so a slow/offline
        # cratescan.service can never block __init__ or the state machine.
        self._gradio_client = None
        self._gradio_lock = threading.Lock()
        if GRADIO_CLIENT_AVAILABLE:
            self._gradio_connect_thread = threading.Thread(target=self._gradio_connect_loop, daemon=True)
            self._gradio_connect_thread.start()

        # Read the physical button state once — BEFORE registering the callback
        # so no interrupt can race against this read.
        self.force_initial_state()

        # No callbacks — button state is hard-polled every 50 ms in run().

        # Grab every GPIO device now; if a pin is still held by a crashed
        # previous instance, the healing loop below keeps retrying instead
        # of crashing __init__ and forcing a manual restart.
        self._acquire_hardware()

        # Sync the RGB LED to the startup state now that rgb_led is available.
        self._set_rgb("purple" if self.state == IDLE else "white")

        self._healing_thread = threading.Thread(target=self._healing_loop, daemon=True)
        self._healing_thread.start()

    # ── startup state sync ───────────────────────────────────────────────────────
    def force_initial_state(self):
        """Read the physical pause-button pin exactly once at startup and sync
        the state machine to physical reality before the main loop starts.
        Uses the same is_pressed convention as _toggle_pause:
          is_pressed=True  → button in RUN  position → WAITING
          is_pressed=False → button in PAUSE position → IDLE
        The 0.1 s settle-time ensures the lgpio sampler is stable."""
        time.sleep(0.1)
        if self.pause_button.is_pressed:
            self.state = WAITING
            print("[INIT] Button in RUN position — starting in WAITING mode.")
        else:
            self.state = IDLE
            print("[INIT] Button in PAUSE position — starting in IDLE mode.")

    # ── hardware acquisition / self-healing  (no manual restarts needed) ───────
    def _acquire_hardware(self):
        with self._lock:
            if self.rgb_led is None:
                try:
                    self.rgb_led = RGBLED(red=RGB_RED_PIN, green=RGB_GREEN_PIN, blue=RGB_BLUE_PIN, active_high=False)
                    print("[HW] RGB LED acquired.")
                except Exception as e:
                    print(f"[HW] RGB LED busy, will retry: {e}")

            if self.led_strips is None:
                try:
                    self.led_strips = Motor(forward=LED_STRIP_IN1, backward=LED_STRIP_IN2)
                    print("[HW] LED strips acquired.")
                except Exception as e:
                    print(f"[HW] LED strips busy, will retry: {e}")

            if self.state == IDLE:
                return    # paused on purpose — don't grab the motor/sensor back from the Debug page

            if self.conveyor is None:
                try:
                    self.conveyor = Motor(forward=CONVEYOR_IN1, backward=CONVEYOR_IN2)
                    print("[HW] Conveyor motor acquired.")
                except Exception as e:
                    print(f"[HW] Conveyor motor busy, will retry: {e}")

            if self.distance_sensor is None:
                try:
                    self.distance_sensor = DistanceSensor(echo=ECHO_PIN, trigger=TRIGGER_PIN, max_distance=2)
                    self._distance_history.clear()
                    print("[HW] Distance sensor acquired.")
                except Exception as e:
                    print(f"[HW] Distance sensor busy, will retry: {e}")

    def _healing_loop(self):
        """Background retry: picks up any GPIO device as soon as it becomes free."""
        while self._running:
            time.sleep(HARDWARE_RETRY_SECONDS)
            if self.rgb_led is None or self.led_strips is None or self.conveyor is None or self.distance_sensor is None:
                self._acquire_hardware()

    def _gradio_connect_loop(self):
        """Keeps retrying the Gradio API connection in the background — used
        only for capture_usb_camera_frame/run_yolo, never for actuators, so
        cratescan.service being offline or slow can never block this process."""
        while self._running:
            if self._gradio_client is None:
                try:
                    client = Client(GRADIO_URL, httpx_kwargs={"timeout": GRADIO_CONNECT_TIMEOUT})
                    with self._gradio_lock:
                        self._gradio_client = client
                    print("[HW] Connected to Gradio API.")
                except Exception as e:
                    print(f"[HW] Gradio API unreachable, will retry: {e}")
            time.sleep(GRADIO_RETRY_SECONDS)

    # ── sensor helper  (moving-average filter over the per-call median) ────────
    def read_distance_cm(self, samples=5):
        """Median of several quick readings, then averaged over the last few
        calls — this ignores both single-sample HC-SR04 noise and brief
        spikes (e.g. the instant a crate leaves the tunnel), so a one-off
        bad reading can't flip the state machine or re-trigger a scan."""
        with self._lock:
            sensor = self.distance_sensor
        if sensor is None:    # released while IDLE, or not yet healed
            return None

        readings = []
        for _ in range(samples):
            try:
                raw = sensor.distance * 100
            except Exception:
                break    # sensor was closed mid-read (e.g. pause fired) — use what we have
            if 2.0 < raw < 190.0:
                readings.append(raw)
            time.sleep(0.01)

        if readings:
            readings.sort()
            self._distance_history.append(readings[len(readings) // 2])

        if not self._distance_history:
            return None
        return sum(self._distance_history) / len(self._distance_history)

    # ── pause / resume  (fires on the GPIO 5 interrupt) ────────────────────────
    def _release_hardware(self):
        """Unconditionally releases the conveyor + sensor pins so the Debug page
        can use them. Each close() is independently guarded — one failing can't
        stop the other from releasing, and can't stop the state machine from
        reaching IDLE. The RGB LED and LED strips are NOT released here — this
        process owns them permanently (see module docstring)."""
        if self.conveyor is not None:
            try:
                self.conveyor.stop()
                self.conveyor.close()
            except Exception as e:
                print(f"[WARN] Error releasing conveyor: {e}")
            self.conveyor = None

        if self.distance_sensor is not None:
            try:
                self.distance_sensor.close()
            except Exception as e:
                print(f"[WARN] Error releasing distance sensor: {e}")
            self.distance_sensor = None

        self._distance_history.clear()

    def _toggle_pause(self):
        now = time.time()
        if now - self._last_toggle_time < PAUSE_DEBOUNCE_SECONDS:
            return
        self._last_toggle_time = now

        # Read the physical pin state — this is the source of truth.
        # is_pressed=True  → button in RUN  position → exit IDLE, go to WAITING.
        # is_pressed=False → button in PAUSE position → enter IDLE.
        button_in_run = self.pause_button.is_pressed

        if button_in_run:
            with self._lock:
                print("[RUN] Button in RUN position — forcing WAITING, conveyor ready.")
                self.state = WAITING
            try:
                self._acquire_hardware()
            except Exception as e:
                print(f"[ERROR] _acquire_hardware() failed unexpectedly: {e}")
            try:
                self._set_rgb("white")
            except Exception as e:
                print(f"[WARN] Could not update RGB LED: {e}")
        else:
            with self._lock:
                print("[PAUSE] Button in PAUSE position — forcing IDLE, conveyor off.")
                self._release_hardware()
                self.state = IDLE
            try:
                self._leds_off()
            except Exception as e:
                print(f"[WARN] Could not turn off LED strips: {e}")
            try:
                self._set_rgb("purple")
            except Exception as e:
                print(f"[WARN] Could not update RGB LED: {e}")

    # ── actuators — direct gpiozero, no network involved ───────────────────────
    def _leds_on(self):
        if self.led_strips is not None:
            self.led_strips.backward(speed=1.0)

    def _leds_off(self):
        if self.led_strips is not None:
            self.led_strips.stop()

    def _set_rgb(self, color):
        if self.rgb_led is not None:
            self.rgb_led.color = RGB_COLORS.get(color, (0, 0, 0))

    def _show_result_color(self, result):
        """Red/yellow/green for RESULT_COLOR_SECONDS based on the scan result."""
        try:
            scan_text = result[1].lower() if isinstance(result, (list, tuple)) and len(result) > 1 else ""
        except Exception:
            scan_text = ""

        # PRIORITEIT 1: Absolute afkeuring (Structural)
        if "structural" in scan_text:
            color = "red"
            
        # PRIORITEIT 2: Vuil aanwezig (Modder, Sticker, Organisch) overschrijft algemene bad_crate
        elif "mud" in scan_text or "sticker" in scan_text or "organic" in scan_text:
            color = "yellow"
            
        # PRIORITEIT 3: YOLO zegt dat het een slechte krat is, maar zag geen specifiek defect
        elif "bad_crate" in scan_text:
            color = "red"
            
        # PRIORITEIT 4: Clean of niks gevonden
        else:
            color = "green"

        self._set_rgb(color)
        threading.Timer(RESULT_COLOR_SECONDS, lambda: self._set_rgb("purple" if self.state == IDLE else "white")).start()
        
    def trigger_scan(self):
        """Calls capture_usb_camera_frame + run_yolo via the Gradio API — data/
        compute only, since the camera + model live in app.py's process."""
        print("[SCAN] Triggering inference...")
        if not GRADIO_CLIENT_AVAILABLE:
            print("[SCAN] (placeholder) gradio_client not installed — simulating inference.")
            return False

        with self._gradio_lock:
            client = self._gradio_client
        if client is None:    # offline/slow right now — _gradio_connect_loop will retry on its own
            print("[SCAN] Gradio app not connected — simulating inference instead.")
            return False

        try:
            frame = client.predict(api_name="/capture_usb_camera_frame")
            print("[SCAN] Frame captured, running YOLO...")
            result = client.predict(
                handle_file(frame), 0.25, 0.45, {},   # conf, iou match the UI slider defaults; {} = fresh counts
                api_name="/run_yolo",
            )
            print(f"[SCAN] run_yolo result: {result}")
            try:
                self._show_result_color(result)
            except Exception as e:
                print(f"[WARN] Could not show result color: {e}")
            return True
        except Exception as e:
            print(f"[SCAN] Gradio API call failed: {e}")
            with self._gradio_lock:
                if self._gradio_client is client:
                    self._gradio_client = None   # force a reconnect before the next scan
            return False

    # ── pausable sleep ──────────────────────────────────────────────────────
    def _interruptible_sleep(self, seconds):
        """Sleep in 50 ms steps; returns early if the main loop toggled us to IDLE."""
        end = time.time() + seconds
        while time.time() < end:
            if self.state == IDLE:
                return
            time.sleep(min(0.05, max(0.0, end - time.time())))

    # ── main loop ───────────────────────────────────────────────────────────
    def run(self):
        print("Crate-Scan Pro automation started. Waiting for a crate...")
        try:
            while self._running:
                # ── Hard-poll the pause button every 50 ms (toggle / flip-flop) ──
                # Only act on a rising edge (False → True) so a single click
                # toggles between IDLE and WAITING without needing to hold.
                current_button_state = self.pause_button.is_pressed
                if current_button_state and not self.last_button_state:
                    if self.state == IDLE:
                        print("[POLL] Button clicked — leaving IDLE, going to WAITING.")
                        self.state = WAITING
                        self._acquire_hardware()
                        self._set_rgb("white")
                    else:
                        print("[POLL] Button clicked — entering IDLE.")
                        self._release_hardware()
                        self._leds_off()
                        self.state = IDLE
                        self._set_rgb("purple")
                self.last_button_state = current_button_state

                # ── Run the state machine only when active ──────────────────────
                if self.state != IDLE:
                    try:
                        self._step()
                    except Exception as e:
                        print(f"[ERROR] State machine step failed, continuing: {e}")

                time.sleep(0.05)   # 50 ms poll cadence

        except KeyboardInterrupt:
            print("\nAutomation stopped by user.")
        finally:
            self.shutdown()

    def _step(self):
        if self.state == IDLE:
            return
        if self.conveyor is None or self.distance_sensor is None:
            return    # still waiting on the healing loop to reacquire a busy pin

        distance = self.read_distance_cm()

        if self.state == WAITING:
            if distance is not None and distance <= ENTRY_DISTANCE_CM:
                print(f"[WAITING] Crate detected at {distance:.1f} cm — starting conveyor.")
                self.state = RUNNING_TO_CAMERA
                self.conveyor.forward()

        elif self.state == RUNNING_TO_CAMERA:
            if distance is not None and distance <= CAMERA_DISTANCE_CM:
                print(f"[RUNNING_TO_CAMERA] Crate under camera at {distance:.1f} cm — stopping conveyor.")
                self.conveyor.stop()
                self.state = SCANNING
                self._leds_on()    # strips on right before the shot, not the whole transit

        elif self.state == SCANNING:
            self._interruptible_sleep(LED_WARMUP_SECONDS)   # minimum settle time under the lights
            if self.state == SCANNING:    # still true unless a pause fired mid warm-up
                self.trigger_scan()
                self._leds_off()          # strips off now that the shot is taken
                if self.conveyor is not None:
                    self.conveyor.forward()   # resume belt to eject the scanned crate
                self._eject_started_at = time.time()
                self.state = EJECTING
            # if paused mid warm-up: motor/sensor already released by _toggle_pause,
            # _previous_state == SCANNING, so resume retries the warm-up from scratch

        elif self.state == EJECTING:
            timed_out = (time.time() - self._eject_started_at) > EJECT_TIMEOUT_SECONDS
            if (distance is not None and distance > ENTRY_DISTANCE_CM) or timed_out:
                print("[EJECTING] Crate cleared the tunnel — waiting for it to fully drop...")
                self.conveyor.stop()
                self._leds_off()
                time.sleep(2.5)   # let a tilting/falling crate fully clear the sensor's field of view
                self._distance_history.clear()   # wipe any residual close-distance readings before WAITING
                self.state = WAITING

    def shutdown(self):
        self._running = False
        self._release_hardware()    # motor + sensor

        if self.led_strips is not None:
            try:
                self.led_strips.stop()
                self.led_strips.close()
            except Exception as e:
                print(f"[WARN] Error releasing LED strips: {e}")
            self.led_strips = None

        if self.rgb_led is not None:
            try:
                self.rgb_led.close()
            except Exception as e:
                print(f"[WARN] Error releasing RGB LED: {e}")
            self.rgb_led = None

        self.pause_button.close()
        print("All GPIO devices released.")


if __name__ == "__main__":
    automation = CrateScanAutomation()
    automation.run()