"""Receive Pupil Capture surface gaze and visualize it over the tag screen."""

from __future__ import annotations

import argparse
from collections import deque
import queue
import statistics
import threading
import time
import tkinter as tk
from dataclasses import dataclass

try:
    import msgpack
    import zmq
except ImportError as exc:
    raise SystemExit(
        "Missing dependency. Run: python -m pip install -r requirements.txt"
    ) from exc

from show_apriltag_surface import (
    SurfaceTagWindow,
    enable_windows_dpi_awareness,
    validate_args,
)


SCREEN_WIDTH = 2560
SCREEN_HEIGHT = 1440


@dataclass(frozen=True)
class GazeSample:
    x_norm: float
    y_norm: float
    confidence: float
    timestamp: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Show screen tags and live surface-relative gaze from Pupil Capture."
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=50020)
    parser.add_argument("--surface", default="screen")
    parser.add_argument("--confidence", type=float, default=0.6)
    parser.add_argument("--dot-size", type=int, default=44)
    parser.add_argument(
        "--smoothing",
        type=float,
        default=0.18,
        help="EMA response from 0 to 1; lower is steadier (default: 0.18)",
    )
    parser.add_argument(
        "--median-window",
        type=int,
        default=5,
        help="number of recent samples used to reject spikes (default: 5)",
    )
    parser.add_argument(
        "--dead-zone",
        type=float,
        default=3.0,
        help="ignore filtered movement smaller than this many pixels (default: 3)",
    )
    parser.add_argument("--tag-ids", nargs=4, type=int, default=(0, 1, 2, 3))
    parser.add_argument("--tag-size", type=int, default=200)
    parser.add_argument("--margin", type=int, default=0)
    parser.add_argument("--background", default="#f2f2f2")
    parser.add_argument("--windowed", action="store_true")
    return parser.parse_args()


class SurfaceReceiver(threading.Thread):
    def __init__(
        self,
        output: queue.Queue,
        stop_event: threading.Event,
        host: str,
        port: int,
        surface: str,
        min_confidence: float,
    ) -> None:
        super().__init__(daemon=True)
        self.output = output
        self.stop_event = stop_event
        self.host = host
        self.port = port
        self.topic = f"surfaces.{surface}"
        self.min_confidence = min_confidence

    def emit(self, item: tuple[str, object]) -> None:
        try:
            self.output.put_nowait(item)
        except queue.Full:
            try:
                self.output.get_nowait()
            except queue.Empty:
                pass
            self.output.put_nowait(item)

    def run(self) -> None:
        context = zmq.Context()
        remote = context.socket(zmq.REQ)
        subscriber = context.socket(zmq.SUB)
        remote.setsockopt(zmq.LINGER, 0)
        subscriber.setsockopt(zmq.LINGER, 0)
        remote.setsockopt(zmq.RCVTIMEO, 3000)
        remote.setsockopt(zmq.SNDTIMEO, 3000)
        try:
            self.emit(("status", f"Connecting to {self.host}:{self.port} …"))
            remote.connect(f"tcp://{self.host}:{self.port}")
            remote.send_string("SUB_PORT")
            sub_port = remote.recv_string()
            subscriber.connect(f"tcp://{self.host}:{sub_port}")
            subscriber.setsockopt_string(zmq.SUBSCRIBE, self.topic)
            poller = zmq.Poller()
            poller.register(subscriber, zmq.POLLIN)
            self.emit(("status", f"Listening for {self.topic}"))

            while not self.stop_event.is_set():
                if subscriber not in dict(poller.poll(100)):
                    continue
                parts = subscriber.recv_multipart()
                if len(parts) < 2:
                    continue
                message = msgpack.unpackb(parts[1], raw=False)
                candidates = (
                    gaze
                    for gaze in message.get("gaze_on_surfaces", ())
                    if gaze.get("on_surf", False)
                    and gaze.get("confidence", 0.0) >= self.min_confidence
                )
                for gaze in candidates:
                    x_norm, y_norm = gaze["norm_pos"]
                    sample = GazeSample(
                        float(x_norm),
                        float(y_norm),
                        float(gaze["confidence"]),
                        float(gaze["timestamp"]),
                    )
                    self.emit(("gaze", sample))
        except zmq.ZMQError as exc:
            self.emit(("error", f"Network error: {exc}"))
        except (KeyError, TypeError, ValueError, msgpack.UnpackException) as exc:
            self.emit(("error", f"Unexpected surface data: {exc}"))
        finally:
            remote.close()
            subscriber.close()
            context.term()


class GazeWindow(SurfaceTagWindow):
    def __init__(self, root: tk.Tk, args: argparse.Namespace) -> None:
        super().__init__(root, args)
        self.args = args
        self.messages: queue.Queue = queue.Queue(maxsize=128)
        self.stop_event = threading.Event()
        self.last_gaze_at = 0.0
        self.x_history: deque[float] = deque(maxlen=args.median_window)
        self.y_history: deque[float] = deque(maxlen=args.median_window)
        self.filtered_position: tuple[float, float] | None = None

        self.status = self.canvas.create_text(
            0,
            18,
            anchor="n",
            fill="#202020",
            font=("Segoe UI", 13, "bold"),
            text="Starting receiver …",
            tags="gaze-ui",
        )
        self.dot = self.canvas.create_oval(
            0, 0, 0, 0, fill="#ff2d2d", outline="white", width=4,
            state="hidden", tags="gaze-ui"
        )
        self.receiver = SurfaceReceiver(
            self.messages,
            self.stop_event,
            args.host,
            args.port,
            args.surface,
            args.confidence,
        )
        self.receiver.start()
        root.protocol("WM_DELETE_WINDOW", self.close)
        root.after(16, self.process_messages)

    def draw(self, event: tk.Event) -> None:
        super().draw(event)
        self.canvas.coords(self.status, event.width / 2, 18)
        self.canvas.tag_raise("gaze-ui")

    def close(self, _event: tk.Event | None = None) -> None:
        self.stop_event.set()
        super().close(_event)

    def process_messages(self) -> None:
        newest: GazeSample | None = None
        try:
            while True:
                kind, value = self.messages.get_nowait()
                if kind == "gaze":
                    newest = value  # type: ignore[assignment]
                else:
                    self.canvas.itemconfigure(self.status, text=str(value))
        except queue.Empty:
            pass

        if newest is not None:
            self.last_gaze_at = time.monotonic()
            raw_x = newest.x_norm * SCREEN_WIDTH
            raw_y = (1.0 - newest.y_norm) * SCREEN_HEIGHT
            self.x_history.append(raw_x)
            self.y_history.append(raw_y)
            median_x = statistics.median(self.x_history)
            median_y = statistics.median(self.y_history)

            if self.filtered_position is None:
                filtered_x, filtered_y = median_x, median_y
            else:
                previous_x, previous_y = self.filtered_position
                alpha = self.args.smoothing
                filtered_x = previous_x + alpha * (median_x - previous_x)
                filtered_y = previous_y + alpha * (median_y - previous_y)
                if abs(filtered_x - previous_x) < self.args.dead_zone:
                    filtered_x = previous_x
                if abs(filtered_y - previous_y) < self.args.dead_zone:
                    filtered_y = previous_y
            self.filtered_position = filtered_x, filtered_y

            x = filtered_x / SCREEN_WIDTH * self.canvas.winfo_width()
            y = filtered_y / SCREEN_HEIGHT * self.canvas.winfo_height()
            radius = self.args.dot_size / 2
            self.canvas.coords(self.dot, x - radius, y - radius, x + radius, y + radius)
            self.canvas.itemconfigure(self.dot, state="normal")
            self.canvas.itemconfigure(
                self.status,
                text=(f"x={filtered_x:7.1f}  y={filtered_y:7.1f} px   "
                      f"confidence={newest.confidence:.2f}"),
            )
            print(
                f"{newest.timestamp:.6f}, {raw_x:.1f}, {raw_y:.1f}, "
                f"{filtered_x:.1f}, {filtered_y:.1f}, {newest.confidence:.3f}",
                flush=True,
            )
        elif time.monotonic() - self.last_gaze_at > 0.35:
            self.canvas.itemconfigure(self.dot, state="hidden")
            self.x_history.clear()
            self.y_history.clear()
            self.filtered_position = None

        if self.root.winfo_exists():
            self.root.after(16, self.process_messages)


def main() -> None:
    args = parse_args()
    validate_args(args)
    if not 0.0 <= args.confidence <= 1.0:
        raise SystemExit("error: --confidence must be between 0 and 1")
    if not 0.0 < args.smoothing <= 1.0:
        raise SystemExit("error: --smoothing must be greater than 0 and at most 1")
    if args.median_window < 1 or args.median_window % 2 == 0:
        raise SystemExit("error: --median-window must be a positive odd number")
    if args.dead_zone < 0:
        raise SystemExit("error: --dead-zone cannot be negative")
    enable_windows_dpi_awareness()
    root = tk.Tk()
    GazeWindow(root, args)
    root.mainloop()


if __name__ == "__main__":
    main()
