# Air Defence System — MicroPython Firmware  (Rev3 Corrected)
# Raspberry Pi Pico  |  Wokwi simulation or real hardware
#
# GPIO mapping (Pico rotated 270° on half-breadboard):
#   GP0  = L293D IN1  (motor dir A)
#   GP1  = L293D IN2  (motor dir B)
#   GP2  = L293D EN1  (motor PWM speed, 1 kHz)
#   GP3  = Servo 1 PWM  (sensor sweep servo)
#   GP4  = Servo 2 PWM  (launch trigger servo)
#   GP5  = HC-SR04 TRIG
#   GP6  = HC-SR04 ECHO  (via 1k+2k voltage divider)
#   GP8  = Buzzer
#   GP9  = Mode switch
#   GP10 = Red LEDs x2
#   GP11 = Green LEDs x2
#   GP13 = Power/white LED
#   GP17 = I2C SDA  (OLED) — GP17 because Pico is rotated, SDA wire lands here
#   GP16 = I2C SCL  (OLED) — GP16 because Pico is rotated, SCL wire lands here
#   GP26 = ADC0  (potentiometer)

from machine import Pin, PWM, ADC, I2C
import time

# ── Try importing ssd1306 ──────────────────────────────────────────
try:
    import ssd1306
    _HAS_SSD1306 = True
except ImportError:
    _HAS_SSD1306 = False
    print("[WARN] ssd1306.py not found.")

# ═══════════════════════════════════════════════════════════════════
#  PIN DEFINITIONS
# ═══════════════════════════════════════════════════════════════════
motor_in1 = Pin(0, Pin.OUT)
motor_in2 = Pin(1, Pin.OUT)
motor_en  = PWM(Pin(2)); motor_en.freq(1000)

servo_scan    = PWM(Pin(3)); servo_scan.freq(50)
servo_trigger = PWM(Pin(4)); servo_trigger.freq(50)

trig = Pin(5, Pin.OUT)
echo = Pin(6, Pin.IN)

buzzer_pwm = PWM(Pin(8))
buzzer_pwm.freq(2730)
buzzer_pwm.duty_u16(0)

mode_sw = Pin(9, Pin.IN, Pin.PULL_DOWN)

led_red   = Pin(10, Pin.OUT)
led_green = Pin(11, Pin.OUT)
led_power = Pin(13, Pin.OUT)

# I2C — NOTE: GP17=SDA, GP16=SCL because the Pico is rotated 270°
# on the breadboard so the OLED wires reach these pins, not GP14/15.
# I2C bus 0 owns GP16/GP17.
try:
    i2c = I2C(0, sda=Pin(17), scl=Pin(16), freq=400000)
    devs = i2c.scan()
    print(f"[I2C] devices: {[hex(d) for d in devs]}")
except Exception as e:
    print(f"[I2C ERROR] {e}")
    i2c = None

pot = ADC(26)

# ═══════════════════════════════════════════════════════════════════
#  OLED INIT
# ═══════════════════════════════════════════════════════════════════
HAS_OLED = False
if _HAS_SSD1306 and i2c is not None:
    try:
        oled = ssd1306.SSD1306_I2C(128, 64, i2c)
        HAS_OLED = True
        print("[OLED] OK")
    except Exception as e:
        print(f"[OLED ERROR] {e}")

# ═══════════════════════════════════════════════════════════════════
#  HELPER FUNCTIONS
# ═══════════════════════════════════════════════════════════════════
def servo_us(pwm_obj, us):
    duty = int(us / 20000.0 * 65535)
    pwm_obj.duty_u16(duty)

def angle_to_us(deg):
    return int(500 + (deg / 180.0) * 2000)

def get_distance_cm():
    trig.value(0); time.sleep_us(2)
    trig.value(1); time.sleep_us(10)
    trig.value(0)
    t0 = time.ticks_us()
    while echo.value() == 0:
        if time.ticks_diff(time.ticks_us(), t0) > 30000:
            return 999
        t0 = time.ticks_us()
    start = time.ticks_us()
    while echo.value() == 1:
        if time.ticks_diff(time.ticks_us(), start) > 30000:
            return 999
    stop = time.ticks_us()
    return time.ticks_diff(stop, start) / 58.0

def oled_show(line1, line2="", line3=""):
    if HAS_OLED:
        oled.fill(0)
        oled.text(str(line1)[:16], 0, 8)
        oled.text(str(line2)[:16], 0, 28)
        oled.text(str(line3)[:16], 0, 48)
        oled.show()
    else:
        print(f"[OLED] {line1} | {line2} | {line3}")

def motor_forward(duty=20000):
    motor_in1.value(1); motor_in2.value(0)
    motor_en.duty_u16(duty)

def motor_stop():
    motor_en.duty_u16(0)
    motor_in1.value(0); motor_in2.value(0)

def buzzer_on():  buzzer_pwm.duty_u16(32768)
def buzzer_off(): buzzer_pwm.duty_u16(0)

def get_threshold_cm():
    raw = pot.read_u16()
    return 10 + int((raw / 65535.0) * 140)

# ═══════════════════════════════════════════════════════════════════
#  SERVO SWEEP
# ═══════════════════════════════════════════════════════════════════
SWEEP_MIN_DEG  = 0
SWEEP_MAX_DEG  = 180
SWEEP_STEP_DEG = 3
STEP_DELAY_MS  = 40

sweep_angle = 90
sweep_dir   = 1

def sweep_step():
    global sweep_angle, sweep_dir
    sweep_angle += sweep_dir * SWEEP_STEP_DEG
    if sweep_angle >= SWEEP_MAX_DEG:
        sweep_angle = SWEEP_MAX_DEG; sweep_dir = -1
    elif sweep_angle <= SWEEP_MIN_DEG:
        sweep_angle = SWEEP_MIN_DEG; sweep_dir = 1
    servo_us(servo_scan, angle_to_us(sweep_angle))
    return sweep_angle

def sweep_home():
    global sweep_angle, sweep_dir
    sweep_angle = 90; sweep_dir = 1
    servo_us(servo_scan, angle_to_us(90))

# ═══════════════════════════════════════════════════════════════════
#  STATE MACHINE
# ═══════════════════════════════════════════════════════════════════
STATE_IDLE=0; STATE_SCANNING=1; STATE_CONFIRM=2; STATE_ALERT=3; STATE_COOLDOWN=4

state         = STATE_IDLE
confirm_count = 0
detected_dist = 0.0
locked_angle  = 90

# Boot
led_power.value(1)
servo_us(servo_scan,    angle_to_us(90))
servo_us(servo_trigger, angle_to_us(10))
motor_stop()
oled_show("AIR DEFENCE", "SYSTEM READY", "Rev3 Servo-Scan")
time.sleep(1)
print("[BOOT] Air Defence System Ready.")

# ═══════════════════════════════════════════════════════════════════
#  MAIN LOOP
# ═══════════════════════════════════════════════════════════════════
while True:

    if state == STATE_IDLE:
        led_green.value(1); led_red.value(0)
        motor_stop()
        oled_show("[ IDLE ]", "System armed", "Mode: Scan+Alert")
        time.sleep_ms(500)
        state = STATE_SCANNING
        print("[STATE] SCANNING")

    elif state == STATE_SCANNING:
        motor_forward(20000)
        current_angle = sweep_step()
        threshold     = get_threshold_cm()
        dist          = get_distance_cm()
        led_green.value(1 if (current_angle % 30) < 15 else 0)
        led_red.value(0)
        oled_show("[ SCANNING ]",
                  f"D:{dist:.0f}cm T:{threshold}",
                  f"A:{current_angle:03d} {'>>>' if sweep_dir > 0 else '<<<'}")
        print(f"[SCAN] angle={current_angle} dist={dist:.1f}cm thresh={threshold}cm")
        if dist <= threshold:
            motor_stop()
            confirm_count = 1; detected_dist = dist; locked_angle = current_angle
            state = STATE_CONFIRM
            print(f"[STATE] CONFIRM at {dist:.1f}cm angle={current_angle}")
        time.sleep_ms(STEP_DELAY_MS)

    elif state == STATE_CONFIRM:
        led_green.value(1)
        threshold = get_threshold_cm()
        dist      = get_distance_cm()
        oled_show("CONFIRMING...", f"D:{dist:.0f}cm", f"({confirm_count}/3)")
        print(f"[CONFIRM] dist={dist:.1f}cm count={confirm_count}/3")
        if dist <= threshold:
            confirm_count += 1; detected_dist = dist
            if confirm_count >= 3:
                state = STATE_ALERT
                print(f"[STATE] ALERT! {detected_dist:.1f}cm angle={locked_angle}")
        else:
            confirm_count = 0; state = STATE_SCANNING
            print("[STATE] back to SCANNING (false alarm)")
        time.sleep_ms(STEP_DELAY_MS)

    elif state == STATE_ALERT:
        led_green.value(0); led_red.value(1)
        motor_stop()
        servo_us(servo_scan,    angle_to_us(locked_angle))
        servo_us(servo_trigger, angle_to_us(locked_angle))
        oled_show("!! DETECTED !!", f"D:{detected_dist:.0f}cm", f"Az:{locked_angle:03d}deg")
        print(f"[ALERT] {detected_dist:.1f}cm angle={locked_angle}")
        buzzer_on(); time.sleep_ms(2000); buzzer_off()
        state = STATE_COOLDOWN
        print("[STATE] COOLDOWN")

    elif state == STATE_COOLDOWN:
        led_red.value(1); led_green.value(0)
        buzzer_off(); motor_stop()
        for i in range(5, 0, -1):
            oled_show("[ COOLDOWN ]", f"Resume in {i}s", "Target Locked")
            print(f"[COOLDOWN] {i}s remaining")
            time.sleep_ms(1000)
        sweep_home()
        servo_us(servo_trigger, angle_to_us(90))
        led_red.value(0)
        confirm_count = 0
        state = STATE_SCANNING
        print("[STATE] SCANNING (after cooldown)")