"""
Aurora Camera System - Pico 2 W Controller
- UART communication with Pi Zero 2 W
- Pan/Tilt servo control via Kitronik Robotics Board
- HC-SR04 ultrasonic sensor with radar display
- HC-SR501 PIR motion sensor with LED indicators
- Auto-scan mode after inactivity timeout
"""

import time
import json
from machine import Pin, SPI, UART
import _thread
import utime
import math

# ============================================================================
# PIN CONFIGURATION
# ============================================================================
# UART Communication
UART_TX_PIN = 0
UART_RX_PIN = 1
UART_BAUDRATE = 115200

# HC-SR04 Ultrasonic Sensor
ULTRASONIC_TRIG_PIN = 14
ULTRASONIC_ECHO_PIN = 15

# PIR Sensor and White LED Headlights
PIR_SENSOR_PIN = 16
WHITE_LED_1_PIN = 17   # White LED headlight 1
WHITE_LED_2_PIN = 19   # White LED headlight 2

# Servo Configuration
SERVO_PAN_CHANNEL = 1
SERVO_TILT_CHANNEL = 2

# ============================================================================
# SYSTEM CONFIGURATION
# ============================================================================
# Servo limits (degrees, servo expects 0-180)
PAN_MIN, PAN_MAX = -90, 90        # Full 180° range
TILT_MIN, TILT_MAX = -40, 45      # Limited: -40° down (to avoid base collision), +45° up

# Auto-scan configuration
INACTIVITY_TIMEOUT = 120  # 2 minutes in seconds
AUTO_SCAN_SPEED = 5       # Degrees per step for scanning (increased from 2)

# Current state
current_pan = 0
current_tilt = 0
last_command_time = 0
auto_scan_mode = False
manual_control_mode = False

# ============================================================================
# KITRONIK ROBOTICS BOARD CLASS
# ============================================================================
class KitronikPicoRobotics:
    SRV_REG_BASE = 0x08
    MOT_REG_BASE = 0x28
    REG_OFFSET = 4
    PRESCALE_VAL = b'\x79'

    def initPCA(self):
        self.i2c.writeto(0, b"\x06")
        self.i2c.writeto_mem(108, 0xfe, self.PRESCALE_VAL)
        self.i2c.writeto_mem(108, 0xfa, b"\x00")
        self.i2c.writeto_mem(108, 0xfb, b"\x00")
        self.i2c.writeto_mem(108, 0xfc, b"\x00")
        self.i2c.writeto_mem(108, 0xfd, b"\x00")
        self.i2c.writeto_mem(108, 0x00, b"\x01")
        utime.sleep_us(500)

    def servoWrite(self, servo, degrees):
        if degrees > 180:
            degrees = 180
        elif degrees < 0:
            degrees = 0
        if servo < 1 or servo > 8:
            raise Exception("INVALID SERVO NUMBER")
        
        calcServo = self.SRV_REG_BASE + ((servo - 1) * self.REG_OFFSET)
        PWMVal = int((degrees * 2.2755) + 102)
        lowByte = PWMVal & 0xFF
        highByte = (PWMVal >> 8) & 0x01
        
        self.i2c.writeto_mem(self.CHIP_ADDRESS, calcServo, bytes([lowByte]))
        self.i2c.writeto_mem(self.CHIP_ADDRESS, calcServo + 1, bytes([highByte]))

    def __init__(self, I2CAddress=108, sda=8, scl=9):
        import machine
        self.CHIP_ADDRESS = 108
        sda = machine.Pin(sda)
        scl = machine.Pin(scl)
        self.i2c = machine.I2C(0, sda=sda, scl=scl, freq=100000)
        self.initPCA()

# ============================================================================
# ST7789 DISPLAY CLASS
# ============================================================================
class ST7789Display:
    def __init__(self, spi, width, height, reset, cs, dc, backlight):
        self.spi = spi
        self.width = height   # Swapped for landscape
        self.height = width
        self.reset = reset
        self.cs = cs
        self.dc = dc
        self.backlight = backlight
        
        # Colors
        self.BLACK = 0x0000
        self.WHITE = 0xFFFF
        self.RED = 0xF800
        self.GREEN = 0x07E0
        self.BLUE = 0x001F
        self.YELLOW = 0xFFE0
        self.CYAN = 0x07FF
        self.MAGENTA = 0xF81F
        
        self.reset.value(1)
        self.cs.value(1)
        self.dc.value(1)
        self.backlight.value(1)
    
    def write_cmd(self, cmd):
        self.cs.value(1)
        self.dc.value(0)
        self.cs.value(0)
        self.spi.write(bytearray([cmd]))
        self.cs.value(1)
    
    def write_data(self, data):
        self.cs.value(1)
        self.dc.value(1)
        self.cs.value(0)
        if isinstance(data, int):
            self.spi.write(bytearray([data]))
        else:
            self.spi.write(data)
        self.cs.value(1)
    
    def init_display(self):
        self.reset.value(1)
        time.sleep_ms(100)
        self.reset.value(0)
        time.sleep_ms(100)
        self.reset.value(1)
        time.sleep_ms(100)
        
        self.write_cmd(0x01)  # Software reset
        time.sleep_ms(150)
        self.write_cmd(0x11)  # Sleep out
        time.sleep_ms(120)
        self.write_cmd(0x3A)  # Color mode
        self.write_data(0x05)  # 16-bit
        self.write_cmd(0x36)  # Memory access
        self.write_data(0x60)  # Rotate 90°
        self.write_cmd(0x2A)  # Column set
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0x01)
        self.write_data(0x3F)
        self.write_cmd(0x2B)  # Row set
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0xEF)
        self.write_cmd(0x21)  # Inversion on
        self.write_cmd(0x29)  # Display on
        time.sleep_ms(100)
    
    def fill_screen(self, color):
        self.write_cmd(0x2A)
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0x01)
        self.write_data(0x3F)
        self.write_cmd(0x2B)
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0x00)
        self.write_data(0xEF)
        self.write_cmd(0x2C)
        
        color_bytes = bytearray([color >> 8, color & 0xFF])
        self.cs.value(1)
        self.dc.value(1)
        self.cs.value(0)
        for _ in range(self.width * self.height):
            self.spi.write(color_bytes)
        self.cs.value(1)
    
    def draw_pixel(self, x, y, color):
        if x < 0 or x >= self.width or y < 0 or y >= self.height:
            return
        self.write_cmd(0x2A)
        self.write_data(x >> 8)
        self.write_data(x & 0xFF)
        self.write_data(x >> 8)
        self.write_data(x & 0xFF)
        self.write_cmd(0x2B)
        self.write_data(y >> 8)
        self.write_data(y & 0xFF)
        self.write_data(y >> 8)
        self.write_data(y & 0xFF)
        self.write_cmd(0x2C)
        self.write_data(color >> 8)
        self.write_data(color & 0xFF)
    
    def draw_line(self, x0, y0, x1, y1, color):
        """Draw line using Bresenham's algorithm"""
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        sx = 1 if x0 < x1 else -1
        sy = 1 if y0 < y1 else -1
        err = dx - dy
        
        while True:
            self.draw_pixel(x0, y0, color)
            if x0 == x1 and y0 == y1:
                break
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x0 += sx
            if e2 < dx:
                err += dx
                y0 += sy
    
    def draw_circle(self, x0, y0, radius, color):
        """Draw circle outline"""
        x = radius
        y = 0
        err = 0
        
        while x >= y:
            self.draw_pixel(x0 + x, y0 + y, color)
            self.draw_pixel(x0 + y, y0 + x, color)
            self.draw_pixel(x0 - y, y0 + x, color)
            self.draw_pixel(x0 - x, y0 + y, color)
            self.draw_pixel(x0 - x, y0 - y, color)
            self.draw_pixel(x0 - y, y0 - x, color)
            self.draw_pixel(x0 + y, y0 - x, color)
            self.draw_pixel(x0 + x, y0 - y, color)
            
            y += 1
            err += 1 + 2 * y
            if 2 * (err - x) + 1 > 0:
                x -= 1
                err += 1 - 2 * x
    
    def show_text_large(self, text, color=None):
        """Show colored screen as text placeholder (no font rendering yet)"""
        if color is None:
            color = self.GREEN
        self.fill_screen(color)
        print(f"Display: {text}")

# ============================================================================
# ULTRASONIC SENSOR CLASS
# ============================================================================
class UltrasonicSensor:
    def __init__(self, trig_pin, echo_pin):
        self.trig = Pin(trig_pin, Pin.OUT)
        self.echo = Pin(echo_pin, Pin.IN)
        self.trig.value(0)
        print(f"🔊 Ultrasonic sensor initialized (Trig=GP{trig_pin}, Echo=GP{echo_pin})")
    
    def measure_distance(self):
        """Measure distance in cm"""
        try:
            self.trig.value(0)
            utime.sleep_us(2)
            self.trig.value(1)
            utime.sleep_us(10)
            self.trig.value(0)
            
            # Timeout after 30ms (corresponds to ~5m range)
            timeout = 30000
            start = utime.ticks_us()
            
            # Wait for echo pulse start
            while self.echo.value() == 0:
                if utime.ticks_diff(utime.ticks_us(), start) > timeout:
                    return -1
            
            pulse_start = utime.ticks_us()
            
            # Wait for echo pulse end
            while self.echo.value() == 1:
                if utime.ticks_diff(utime.ticks_us(), pulse_start) > timeout:
                    return -1
            
            pulse_end = utime.ticks_us()
            pulse_duration = utime.ticks_diff(pulse_end, pulse_start)
            
            # Calculate distance (speed of sound = 343 m/s = 0.0343 cm/μs)
            distance = (pulse_duration * 0.0343) / 2
            
            return distance if distance < 400 else -1  # Max 4m range
            
        except:
            return -1

# ============================================================================
# MOTION DETECTOR CLASS (with White LED Headlights)
# ============================================================================
class MotionDetector:
    def __init__(self, pir_pin, led1_pin, led2_pin):
        self.pir = Pin(pir_pin, Pin.IN, Pin.PULL_DOWN)
        self.led1 = Pin(led1_pin, Pin.OUT)
        self.led2 = Pin(led2_pin, Pin.OUT)
        self.motion_detected = False
        self.motion_count = 0
        self.running = False
        
        # Turn off both white LEDs initially
        self.led1.value(0)
        self.led2.value(0)
        
        print(f"🔍 Motion detector initialized")
        print(f"   PIR: GP{pir_pin}")
        print(f"   White LED 1: GP{led1_pin}, White LED 2: GP{led2_pin}")
    
    def check_motion(self):
        return self.pir.value() == 1
    
    def turn_on_headlights(self):
        """Turn on both white LED headlights"""
        self.led1.value(1)
        self.led2.value(1)
    
    def turn_off_headlights(self):
        """Turn off both white LED headlights"""
        self.led1.value(0)
        self.led2.value(0)
    
    def monitor_loop(self):
        """Background thread for motion monitoring"""
        print("👁️  Motion monitoring started")
        self.running = True
        
        while self.running:
            if self.check_motion():
                if not self.motion_detected:
                    self.motion_detected = True
                    self.motion_count += 1
                    self.turn_on_headlights()
                    print(f"🚨 Motion detected! Headlights ON (Count: {self.motion_count})")
            else:
                if self.motion_detected:
                    self.motion_detected = False
                    self.turn_off_headlights()
                    print("✓ Motion ended - Headlights OFF")
            
            time.sleep_ms(100)  # Check every 100ms
    
    def stop(self):
        self.running = False
        self.turn_off_headlights()

# ============================================================================
# SERVO CONTROLLER CLASS
# ============================================================================
class ServoController:
    def __init__(self, pan_channel, tilt_channel):
        self.robotics_board = KitronikPicoRobotics()
        self.pan_channel = pan_channel
        self.tilt_channel = tilt_channel
        print("🤖 Kitronik Robotics Board initialized")
    
    def angle_to_servo(self, angle):
        """Convert -90/+90 angle to 0-180 servo range"""
        return int(angle + 90)
    
    def set_position(self, pan, tilt):
        """Set servo positions with limits"""
        global current_pan, current_tilt
        
        pan = max(PAN_MIN, min(PAN_MAX, pan))
        tilt = max(TILT_MIN, min(TILT_MAX, tilt))
        
        pan_servo = self.angle_to_servo(pan)
        tilt_servo = self.angle_to_servo(tilt)
        
        try:
            self.robotics_board.servoWrite(self.pan_channel, pan_servo)
            self.robotics_board.servoWrite(self.tilt_channel, tilt_servo)
            current_pan = pan
            current_tilt = tilt
            return True
        except Exception as e:
            print(f"❌ Servo error: {e}")
            return False
    
    def smooth_move(self, target_pan, target_tilt, steps=10):
        """Smooth movement to target"""
        start_pan, start_tilt = current_pan, current_tilt
        
        for i in range(steps + 1):
            progress = i / steps
            pan = start_pan + (target_pan - start_pan) * progress
            tilt = start_tilt + (target_tilt - start_tilt) * progress
            self.set_position(pan, tilt)
            time.sleep_ms(20)
    
    def startup_sequence(self):
        """Startup servo sweep"""
        print("🎯 Running startup sequence...")
        
        # Pan sweep: -90 → 0 → +90 → 0
        positions = [
            (-90, 0), (0, 0), (90, 0), (0, 0),
            (0, -45), (0, 0), (0, 45), (0, 0)
        ]
        
        for pan, tilt in positions:
            self.smooth_move(pan, tilt, steps=15)
            time.sleep_ms(300)
        
        print("✅ Startup sequence complete")

# ============================================================================
# UART HANDLER CLASS
# ============================================================================
class UARTHandler:
    def __init__(self, tx_pin, rx_pin, baudrate):
        self.uart = UART(0, baudrate=baudrate, tx=Pin(tx_pin), rx=Pin(rx_pin))
        self.rx_buffer = ""
        print(f"📡 UART initialized: {baudrate} baud (TX=GP{tx_pin}, RX=GP{rx_pin})")
    
    def send_response(self, data):
        """Send JSON response"""
        msg = json.dumps(data) + "\n"
        self.uart.write(msg.encode())
    
    def send_event(self, event, data):
        """Send event message"""
        msg = {"type": "event", "event": event}
        msg.update(data)
        self.send_response(msg)
    
    def check_messages(self):
        """Check for incoming commands"""
        if self.uart.any():
            data = self.uart.read()
            if data:
                self.rx_buffer += data.decode('utf-8', 'ignore')
                
                # Process complete messages (terminated by \n)
                while '\n' in self.rx_buffer:
                    line, self.rx_buffer = self.rx_buffer.split('\n', 1)
                    if line.strip():
                        try:
                            return json.loads(line)
                        except:
                            print(f"❌ Invalid JSON: {line}")
        return None

# ============================================================================
# RADAR DISPLAY MODE
# ============================================================================
def draw_radar_display(display, angle, distance):
    """Draw mini radar screen"""
    # Center and radius
    cx, cy = 160, 180
    max_r = 100
    
    # Don't clear screen - keep previous dots for persistence
    # Only redraw circles and scan line
    
    # Draw radar circles
    display.draw_circle(cx, cy, 25, display.GREEN)
    display.draw_circle(cx, cy, 50, display.GREEN)
    display.draw_circle(cx, cy, 75, display.GREEN)
    display.draw_circle(cx, cy, 100, display.GREEN)
    
    # Draw scan line
    rad = math.radians(angle)
    end_x = int(cx + max_r * math.cos(rad))
    end_y = int(cy - max_r * math.sin(rad))
    display.draw_line(cx, cy, end_x, end_y, display.YELLOW)
    
    # Draw distance point - make it bigger and brighter
    if 0 < distance < 100:
        dist_r = int(distance * max_r / 100)
        point_x = int(cx + dist_r * math.cos(rad))
        point_y = int(cy - dist_r * math.sin(rad))
        # Draw a small circle instead of single pixel for visibility
        for dx in range(-2, 3):
            for dy in range(-2, 3):
                if dx*dx + dy*dy <= 4:  # Circle of radius 2
                    display.draw_pixel(point_x + dx, point_y + dy, display.RED)
        print(f"📏 Object detected at {distance:.1f}cm, angle {angle}°")

# ============================================================================
# MAIN LOOP
# ============================================================================
def main():
    global last_command_time, auto_scan_mode, manual_control_mode
    
    print("🌅 Aurora Camera System Starting...")
    
    # Initialize hardware
    onboard_led = Pin("LED", Pin.OUT)
    onboard_led.value(1)  # Turn on onboard LED to indicate power
    time.sleep(2)
    
    # Initialize display
    print("📱 Initializing display...")
    spi = SPI(0, baudrate=40000000, sck=Pin(6), mosi=Pin(7))
    display = ST7789Display(
        spi=spi, width=240, height=320,
        reset=Pin(3, Pin.OUT), cs=Pin(5, Pin.OUT),
        dc=Pin(4, Pin.OUT), backlight=Pin(2, Pin.OUT)
    )
    display.init_display()
    display.fill_screen(display.BLUE)  # Blue screen = Aurora starting
    print("Display: AURORA")
    time.sleep(1)
    display.fill_screen(display.YELLOW)  # Yellow screen = Loading
    print("Display: LOADING...")
    
    # Initialize servo controller
    print("🔧 Initializing servos...")
    servo = ServoController(SERVO_PAN_CHANNEL, SERVO_TILT_CHANNEL)
    
    # Run startup sequence
    servo.startup_sequence()
    
    # Initialize ultrasonic sensor
    print("🔊 Initializing ultrasonic sensor...")
    ultrasonic = UltrasonicSensor(ULTRASONIC_TRIG_PIN, ULTRASONIC_ECHO_PIN)
    
    # Initialize motion detector
    print("🔍 Initializing motion detector...")
    motion = MotionDetector(PIR_SENSOR_PIN, WHITE_LED_1_PIN, WHITE_LED_2_PIN)
    _thread.start_new_thread(motion.monitor_loop, ())
    
    # Initialize UART
    print("📡 Initializing UART...")
    uart = UARTHandler(UART_TX_PIN, UART_RX_PIN, UART_BAUDRATE)
    
    display.fill_screen(display.GREEN)  # Green screen = Ready
    print("Display: READY")
    time.sleep(1)
    display.fill_screen(display.BLACK)
    
    print("✅ Aurora system ready!")
    print("👁️  Waiting for UART commands...")
    
    last_command_time = time.time()
    scan_angle = 0
    scan_direction = 1
    
    # Main loop
    while True:
        current_time = time.time()
        
        # Check for UART commands
        cmd = uart.check_messages()
        if cmd:
            last_command_time = current_time
            auto_scan_mode = False
            manual_control_mode = True
            
            cmd_type = cmd.get('type')
            if cmd_type == 'cmd':
                command = cmd.get('cmd')
                
                if command == 'move':
                    pan = cmd.get('pan', current_pan)
                    tilt = cmd.get('tilt', current_tilt)
                    smooth = cmd.get('smooth', False)
                    
                    if smooth:
                        servo.smooth_move(pan, tilt)
                    else:
                        servo.set_position(pan, tilt)
                    
                    uart.send_response({
                        'type': 'resp',
                        'status': 'ok',
                        'pan': current_pan,
                        'tilt': current_tilt
                    })
                    
                    display.fill_screen(display.CYAN)  # Cyan = Manual Control
                    print("Display: MANUAL CONTROL")
                
                elif command == 'center':
                    servo.smooth_move(0, 0)
                    uart.send_response({'type': 'resp', 'status': 'ok'})
                
                elif command == 'get_distance':
                    dist = ultrasonic.measure_distance()
                    uart.send_response({
                        'type': 'resp',
                        'status': 'ok',
                        'distance': dist
                    })
        
        # Check inactivity timeout
        if current_time - last_command_time > INACTIVITY_TIMEOUT:
            if not auto_scan_mode:
                print("⏰ Inactivity timeout - starting auto-scan mode")
                auto_scan_mode = True
                manual_control_mode = False
                display.fill_screen(display.BLACK)
            
            # Auto-scan with radar display
            servo.set_position(scan_angle - 90, 0)
            distance = ultrasonic.measure_distance()
            
            draw_radar_display(display, scan_angle, distance)
            
            # Send distance event
            uart.send_event('distance', {
                'angle': scan_angle,
                'distance': distance
            })
            
            scan_angle += scan_direction * AUTO_SCAN_SPEED
            if scan_angle >= 180:
                scan_angle = 180
                scan_direction = -1
            elif scan_angle <= 0:
                scan_angle = 0
                scan_direction = 1
        
        time.sleep_ms(50)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n🛑 Aurora stopped")
    except Exception as e:
        print(f"💥 Error: {e}")