#!/usr/bin/env python3
import sys
import os
import signal

# Cleanly handle parent pipe closure on SIGTERM / exit without throwing BrokenPipeError
if hasattr(signal, 'SIGPIPE'):
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)

os.environ["OPENCV_LOG_LEVEL"] = "SILENT"
import json
import time
import base64
import threading
import subprocess
import cv2
import mss
import mss.tools
import numpy as np

# We capture a 128x128 full-word region around the gaze for whole-word reading
REGION_SIZE = 128
DOWNSAMPLE_SIZE = 128

cursor_x = 1920 // 2
cursor_y = 1080 // 2
mouse_x = 1920 // 2
mouse_y = 1080 // 2
screen_w = 1920
screen_h = 1080
frame_size = 1920 * 1080 * 3

# Neuromorphic V1/V2 Multi-Scale Gabor Wavelet Bank (0°, 45°, 90°, 135°)
GABOR_ORIENTATIONS = [0.0, np.pi / 4.0, np.pi / 2.0, 3.0 * np.pi / 4.0]
GABOR_SCALES = [4.0, 8.0] # High frequency (fine serifs) & Medium frequency (word silhouettes)
gabor_filters = []
for theta in GABOR_ORIENTATIONS:
    for lam in GABOR_SCALES:
        ksize = 9
        sigma = 2.0
        gamma = 0.5
        # Even (Cosine/Real) and Odd (Sine/Imag) quadrature pairs for Phase Congruency
        even_k = cv2.getGaborKernel((ksize, ksize), sigma, theta, lam, gamma, 0.0, ktype=cv2.CV_32F)
        odd_k = cv2.getGaborKernel((ksize, ksize), sigma, theta, lam, gamma, np.pi / 2.0, ktype=cv2.CV_32F)
        gabor_filters.append((even_k, odd_k, theta, lam))

running = True
prev_peri = None

def cleanup_wf():
    global wf_proc
    if wf_proc:
        try:
            wf_proc.terminate()
            wf_proc.wait(timeout=0.2)
        except Exception:
            try:
                wf_proc.kill()
            except Exception:
                pass
        wf_proc = None

# Thread-Safe Real-Time Frame Ingestion Double-Buffer
latest_raw_frame = None
raw_frame_lock = threading.Lock()
frame_ready_event = threading.Event()

def stdin_listener():
    global cursor_x, cursor_y, mouse_x, mouse_y, screen_w, screen_h, running
    while running:
        line = sys.stdin.readline()
        if not line:
            running = False
            break
        try:
            cmd = json.loads(line)
            if cmd.get('type') == 'update_gaze':
                cursor_x = cmd.get('x', cursor_x)
                cursor_y = cmd.get('y', cursor_y)
                mouse_x = cmd.get('mouseX', mouse_x)
                mouse_y = cmd.get('mouseY', mouse_y)
            elif cmd.get('type') == 'update_screen':
                screen_w = cmd.get('w', screen_w)
                screen_h = cmd.get('h', screen_h)
        except Exception:
            pass

def main():
    global screen_w, screen_h, frame_size, running, prev_peri, latest_raw_frame
    print('{"status": "ready"}', flush=True)
    
    t_stdin = threading.Thread(target=stdin_listener, daemon=True)
    t_stdin.start()
    
    capture_mode = "unknown"
    sct = None
    monitor = None
    wf_proc = None

    is_wayland = bool(os.environ.get('WAYLAND_DISPLAY'))

    # 1. Wayland Native Detection
    if is_wayland:
        try:
            res = subprocess.run(["grim", "-t", "ppm", "-"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=1.0)
            if res.returncode == 0 and res.stdout:
                lines = res.stdout.split(b"\n", 3)
                if len(lines) >= 3 and lines[0].startswith(b"P6"):
                    dims = lines[1].strip().split(b" ")
                    if len(dims) == 2:
                        screen_w = int(dims[0])
                        screen_h = int(dims[1])
                        frame_size = screen_w * screen_h * 3
        except Exception:
            pass

        # Try wf-recorder rawvideo RGB24 pipe
        try:
            cmd = ["wf-recorder", "-r", "30", "--muxer=rawvideo", "--codec=rawvideo", "--pixel-format=rgb24", "-f", "pipe:1"]
            wf_proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
            time.sleep(0.05)
            if wf_proc.poll() is None:
                capture_mode = "wayland_wf"
                print(f'{{"status": "info", "message": "Using Wayland (wf-recorder) for real-time desktop capture ({screen_w}x{screen_h})"}}', file=sys.stderr, flush=True)
            else:
                wf_proc = None
        except Exception:
            wf_proc = None

        if capture_mode == "unknown":
            try:
                res = subprocess.run(["grim", "-t", "ppm", "-"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=1.0)
                if res.returncode == 0 and res.stdout:
                    capture_mode = "wayland_grim"
                    print(f'{{"status": "info", "message": "Using Wayland (grim) for screen capture ({screen_w}x{screen_h})"}}', file=sys.stderr, flush=True)
            except Exception:
                pass

    # 2. Native OS / X11 Direct Capture via mss
    if capture_mode == "unknown":
        try:
            sct = mss.MSS()
            if len(sct.monitors) > 0:
                monitor = sct.monitors[1] if len(sct.monitors) > 1 else sct.monitors[0]
                screen_w = monitor["width"]
                screen_h = monitor["height"]
                frame_size = screen_w * screen_h * 3
                capture_mode = "mss"
                print(f'{{"status": "info", "message": "Using X11/Direct (mss) for vision screen capture ({screen_w}x{screen_h})"}}', file=sys.stderr, flush=True)
        except Exception:
            sct = None

    # 3. Fallback: Virtual Headless Frame Generator
    if capture_mode == "unknown":
        capture_mode = "synthetic_fallback"
        print(f'{{"status": "warn", "message": "No active display server found. Using synthetic headless vision buffer ({screen_w}x{screen_h})"}}', file=sys.stderr, flush=True)

    def read_exact_bytes(stream, num_bytes):
        chunks = []
        bytes_read = 0
        while bytes_read < num_bytes:
            chunk = stream.read(num_bytes - bytes_read)
            if not chunk:
                return None
            chunks.append(chunk)
            bytes_read += len(chunk)
        return b''.join(chunks)

    def cleanup_wf():
        nonlocal wf_proc
        if wf_proc:
            try:
                wf_proc.terminate()
                wf_proc.wait(timeout=1.0)
            except Exception:
                try:
                    wf_proc.kill()
                except Exception:
                    pass
            wf_proc = None

    import atexit
    import signal
    atexit.register(cleanup_wf)
    signal.signal(signal.SIGINT, lambda s, f: (cleanup_wf(), sys.exit(0)))
    signal.signal(signal.SIGTERM, lambda s, f: (cleanup_wf(), sys.exit(0)))

    # Continuous Background Frame Reader Thread
    # Drains wf-recorder / mss / grim into latest_raw_frame with ZERO pipe backlog
    def screen_capture_worker():
        nonlocal wf_proc, capture_mode, sct, monitor
        global latest_raw_frame, running, screen_w, screen_h, frame_size
        while running:
            try:
                frame = None
                if capture_mode == "wayland_wf" and wf_proc:
                    if wf_proc.poll() is not None:
                        cleanup_wf()
                        capture_mode = "wayland_grim"
                    else:
                        # Non-blocking OS pipe drain: check if more than 1 frame is queued in pipe buffer
                        if sys.platform.startswith('linux'):
                            try:
                                import fcntl, termios, array
                                buf = array.array('i', [0])
                                fcntl.ioctl(wf_proc.stdout.fileno(), termios.FIONREAD, buf)
                                avail = buf[0]
                                if avail >= frame_size * 2:
                                    drop_bytes = (avail // frame_size - 1) * frame_size
                                    wf_proc.stdout.read(drop_bytes)
                            except Exception:
                                pass

                        raw = read_exact_bytes(wf_proc.stdout, frame_size)
                        if not raw or len(raw) != frame_size:
                            cleanup_wf()
                            capture_mode = "wayland_grim"
                        else:
                            frame = np.frombuffer(raw, dtype=np.uint8).reshape((screen_h, screen_w, 3))

                elif capture_mode == "wayland_grim":
                    res = subprocess.run(["grim", "-t", "ppm", "-"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=0.2)
                    if res.returncode == 0 and res.stdout:
                        bgr_img = cv2.imdecode(np.frombuffer(res.stdout, dtype=np.uint8), cv2.IMREAD_COLOR)
                        if bgr_img is not None:
                            frame = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
                    time.sleep(0.02)

                elif capture_mode == "mss" and sct and monitor:
                    full_bbox = {
                        "top": monitor["top"],
                        "left": monitor["left"],
                        "width": monitor["width"],
                        "height": monitor["height"]
                    }
                    full_sct = sct.grab(full_bbox)
                    frame = np.array(full_sct)[..., :3][..., ::-1]
                    time.sleep(0.02)

                elif capture_mode == "synthetic_fallback":
                    frame = np.zeros((screen_h, screen_w, 3), dtype=np.uint8)
                    t_step = int((time.time() * 30) % 255)
                    frame[100:300, 100:500] = [t_step, 120, 200]
                    time.sleep(1.0 / 30.0)

                if frame is not None:
                    with raw_frame_lock:
                        latest_raw_frame = frame
                    frame_ready_event.set()
            except Exception:
                time.sleep(0.05)

    t_capture = threading.Thread(target=screen_capture_worker, daemon=True)
    t_capture.start()

    # Main Vision Processing Loop (30 FPS, guaranteed real-time)
    while running:
        start_time = time.time()
        
        try:
            full_np = None
            with raw_frame_lock:
                if latest_raw_frame is not None:
                    full_np = latest_raw_frame.copy()

            if full_np is None:
                frame_ready_event.wait(timeout=0.03)
                with raw_frame_lock:
                    if latest_raw_frame is not None:
                        full_np = latest_raw_frame.copy()

            if full_np is None:
                time.sleep(0.01)
                continue
            
            # Physiological Microsaccades (Biological ocular micro-drift / nystagmus)
            micro_jitter_x = int(np.sin(time.time() * 62.83) * 1.5)
            micro_jitter_y = int(np.cos(time.time() * 50.26) * 1.5)
            eff_gaze_x = cursor_x + micro_jitter_x
            eff_gaze_y = cursor_y + micro_jitter_y

            # Fovea region crop from full_np
            left = int(max(0, min(full_np.shape[1] - REGION_SIZE, eff_gaze_x - REGION_SIZE // 2)))
            top = int(max(0, min(full_np.shape[0] - REGION_SIZE, eff_gaze_y - REGION_SIZE // 2)))
            fovea_np = full_np[top:top+REGION_SIZE, left:left+REGION_SIZE].copy()
            
            # Draw system cursor on full_np (peripheral) before downsampling
            sys_cursor = (int(mouse_x), int(mouse_y))
            cv2.line(full_np, (sys_cursor[0]-10, sys_cursor[1]), (sys_cursor[0]+10, sys_cursor[1]), (255, 0, 0), 3)
            cv2.line(full_np, (sys_cursor[0], sys_cursor[1]-10), (sys_cursor[0], sys_cursor[1]+10), (255, 0, 0), 3)
            cv2.circle(full_np, sys_cursor, 4, (0, 0, 255), -1)
            
            # Periphery high-resolution downsample (16:9 ratio, 960x540)
            peri_downsampled = cv2.resize(full_np, (960, 540), interpolation=cv2.INTER_AREA)
            
            # Efficient peripheral retinal encoding (320x180 = 57,600 pixels)
            peri_low = cv2.resize(peri_downsampled, (320, 180), interpolation=cv2.INTER_AREA)
            cached_peri_bytes = peri_low.tobytes()
            
            # Motion detection (Peripheral Saccade Reflex)
            motion_x, motion_y, motion_mag = 0, 0, 0
            peri_gray = cv2.cvtColor(peri_downsampled, cv2.COLOR_RGB2GRAY)
            diff = np.zeros_like(peri_gray)
            if prev_peri is not None:
                diff = cv2.absdiff(peri_gray, prev_peri)
                _, thresh = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY)
                M = cv2.moments(thresh)
                if M["m00"] > 1000:
                    cx = int(M["m10"] / M["m00"])
                    cy = int(M["m01"] / M["m00"])
                    motion_mag = int(M["m00"])
                    motion_x = int((cx / 960.0) * screen_w)
                    motion_y = int((cy / 540.0) * screen_h)
                    cv2.circle(peri_downsampled, (cx, cy), 15, (255, 255, 0), 2)
            prev_peri = peri_gray
            
            # Fovea Patch (128x128 RGB)
            fovea_downsampled = cv2.resize(fovea_np, (DOWNSAMPLE_SIZE, DOWNSAMPLE_SIZE), interpolation=cv2.INTER_LINEAR)
            fovea_normalized = fovea_downsampled / 255.0
            fovea_bytes = fovea_normalized.astype(np.float32).tobytes()
            
            # Base64 encodings
            fovea_b64 = base64.b64encode(fovea_bytes).decode('utf-8')
            peri_b64 = base64.b64encode(cached_peri_bytes).decode('utf-8')
            
            # JPEG encodings for dashboard UI
            fovea_bgr = cv2.cvtColor(fovea_np, cv2.COLOR_RGB2BGR)
            peri_bgr = cv2.cvtColor(peri_downsampled, cv2.COLOR_RGB2BGR)
            _, fovea_jpg = cv2.imencode('.jpg', fovea_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 35])
            _, peri_jpg = cv2.imencode('.jpg', peri_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 35])
            fovea_jpg_b64 = base64.b64encode(fovea_jpg.tobytes()).decode('utf-8')
            peri_jpg_b64 = base64.b64encode(peri_jpg.tobytes()).decode('utf-8')
            
            # Multi-Scale Log-Gabor Orientation Columns
            peri_low_gray = cv2.cvtColor(peri_low, cv2.COLOR_RGB2GRAY)
            peri_gray_f = peri_low_gray.astype(np.float32) / 255.0
            local_energy_sum = np.zeros_like(peri_gray_f)

            for even_k, odd_k, _, _ in gabor_filters:
                r_even = cv2.filter2D(peri_gray_f, cv2.CV_32F, even_k)
                r_odd = cv2.filter2D(peri_gray_f, cv2.CV_32F, odd_k)
                energy = np.sqrt(r_even**2 + r_odd**2)
                local_energy_sum += energy

            phase_congruency = cv2.normalize(local_energy_sum, None, 0.0, 1.0, cv2.NORM_MINMAX)

            # Multi-Class Visual Region Discrimination (Text vs Icons vs Images)
            edge_thresh = (phase_congruency * 255.0).astype(np.uint8)
            _, text_thresh = cv2.threshold(edge_thresh, 35, 255, cv2.THRESH_BINARY)
            contours, _ = cv2.findContours(text_thresh, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)

            text_boxes = []
            icon_boxes = []
            image_boxes = []

            if len(contours) > 0:
                for cnt in contours:
                    x_b, y_b, w_b, h_b = cv2.boundingRect(cnt)
                    area = w_b * h_b
                    if area < 30:
                        continue
                    
                    aspect_ratio = float(w_b) / max(1.0, float(h_b))
                    sx = int((x_b / 320.0) * screen_w)
                    sy = int((y_b / 180.0) * screen_h)
                    sw = int((w_b / 320.0) * screen_w)
                    sh = int((h_b / 180.0) * screen_h)
                    box = [sx, sy, sw, sh]

                    # 1. Large Image / Media / Canvas Region
                    if (w_b >= 80 and h_b >= 50) or area >= 5000:
                        image_boxes.append(box)
                    # 2. Text / Words / Graphemes (horizontal strokes or compact letter blocks)
                    elif (1.2 <= aspect_ratio <= 14.0 and h_b <= 45) or (0.25 <= aspect_ratio <= 1.2 and 6 <= h_b <= 35 and w_b <= 35):
                        text_boxes.append(box)
                    # 3. Icons / UI Buttons / Badges (compact squares, circles, pill buttons)
                    elif (0.7 <= aspect_ratio <= 1.4 and 12 <= w_b <= 75 and 12 <= h_b <= 75) or (1.4 <= aspect_ratio <= 4.0 and 15 <= h_b <= 50):
                        icon_boxes.append(box)
                    else:
                        if h_b <= 30:
                            text_boxes.append(box)
                        else:
                            icon_boxes.append(box)

            text_boxes.sort(key=lambda b: (b[1] // 40, b[0]))
            icon_boxes.sort(key=lambda b: (b[1] // 40, b[0]))
            image_boxes.sort(key=lambda b: (b[1] // 40, b[0]))
            text_boxes = text_boxes[:20]
            icon_boxes = icon_boxes[:15]
            image_boxes = image_boxes[:10]

            msg = {
                "type": "vision_pixels",
                "fovea_b64": fovea_b64,
                "peri_b64": peri_b64,
                "fovea_jpg_b64": fovea_jpg_b64,
                "peri_jpg_b64": peri_jpg_b64,
                "x": left,
                "y": top,
                "gaze_x": cursor_x,
                "gaze_y": cursor_y,
                "motion_x": motion_x,
                "motion_y": motion_y,
                "motion_mag": motion_mag,
                "text_boxes": text_boxes,
                "icon_boxes": icon_boxes,
                "image_boxes": image_boxes
            }
            
            print(json.dumps(msg), flush=True)
                
        except (BrokenPipeError, IOError, OSError):
            sys.exit(0)
        except Exception as e:
            try:
                print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
            except Exception:
                pass
            time.sleep(0.05)
            
        elapsed = time.time() - start_time
        sleep_time = (1.0 / 30.0) - elapsed
        if sleep_time > 0:
            time.sleep(sleep_time)

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)
