# sites/ai_mouse/ai_mouse/_collector.py
from __future__ import annotations

import json
import math
import random
from enum import Enum, auto
from pathlib import Path


class CollectorState(Enum):
    IDLE = auto()
    HOVER_A = auto()
    RECORDING = auto()


class Collector:
    """Manages A→B mouse movement trace collection state and persistence.

    The state-machine methods (_on_mouse_motion, _on_mouseup, _on_skip)
    are public so they can be called by the web API without a display.
    A/B positions are exposed as .a_pos / .b_pos attributes.
    """

    POINT_RADIUS = 15   # pixels — must be inside this to hover/click
    DWELL_MS = 200      # milliseconds to dwell inside A before recording starts

    def __init__(
        self,
        count: int,
        dist_min: int,
        dist_max: int,
        output_path: Path,
        screen_size: tuple[int, int] = (800, 600),
    ):
        self.count = count
        self.dist_min = dist_min
        self.dist_max = dist_max
        self.output_path = Path(output_path)
        self.screen_w, self.screen_h = screen_size

        self.collected = 0
        self.state = CollectorState.IDLE
        self._buffer: list[dict] = []
        self._hover_enter_t: int = 0
        self._record_start_t: int = 0

        self.a_pos, self.b_pos = self._new_ab()

    # ------------------------------------------------------------------
    # State machine (called by web API)
    # ------------------------------------------------------------------

    def _on_mouse_motion(self, mx: int, my: int, t: int) -> None:
        """Handle a MOUSEMOTION event at pixel (mx, my), time t ms from start."""
        if self.state == CollectorState.IDLE:
            if self._inside(mx, my, self.a_pos):
                self.state = CollectorState.HOVER_A
                self._hover_enter_t = t

        elif self.state == CollectorState.HOVER_A:
            if not self._inside(mx, my, self.a_pos):
                self.state = CollectorState.IDLE
            elif t - self._hover_enter_t >= self.DWELL_MS:
                self.state = CollectorState.RECORDING
                self._record_start_t = t
                self._buffer = [{"type": "move", "x": mx, "y": my, "t": 0}]

        elif self.state == CollectorState.RECORDING:
            rel_t = t - self._record_start_t
            self._buffer.append({"type": "move", "x": mx, "y": my, "t": rel_t})

    def _on_mousedown(self, mx: int, my: int, t: int) -> None:
        """Handle a MOUSEBUTTONDOWN event."""
        if self.state == CollectorState.RECORDING:
            rel_t = t - self._record_start_t
            self._buffer.append({"type": "down", "x": mx, "y": my, "t": rel_t})

    def _on_mouseup(self, mx: int, my: int, t: int) -> None:
        """Handle a MOUSEBUTTONUP event."""
        if self.state == CollectorState.RECORDING:
            rel_t = t - self._record_start_t
            self._buffer.append({"type": "up", "x": mx, "y": my, "t": rel_t})

            if self._inside(mx, my, self.b_pos):
                self._save_trace()
                self.collected += 1
                if self.collected < self.count:
                    self.a_pos, self.b_pos = self._new_ab()
                self.state = CollectorState.IDLE
            else:
                # Click outside B — discard buffer and regenerate
                self._on_skip()

    def _on_skip(self) -> None:
        """Handle ESC/skip — discard current buffer, regenerate A/B."""
        self._buffer = []
        self.state = CollectorState.IDLE
        self.a_pos, self.b_pos = self._new_ab()

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def _save_trace(self) -> None:
        dist = self._dist(self.a_pos, self.b_pos)
        angle = math.degrees(
            math.atan2(
                self.b_pos[1] - self.a_pos[1],
                self.b_pos[0] - self.a_pos[0],
            )
        )
        trace = {
            "meta": {
                "start": list(self.a_pos),
                "end": list(self.b_pos),
                "dist": round(dist),
                "angle": round(angle, 1),
            },
            "events": list(self._buffer),
        }
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        with self.output_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(trace, ensure_ascii=False) + "\n")
        self._buffer = []

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _inside(self, mx: int, my: int, pos: tuple[int, int]) -> bool:
        return self._dist((mx, my), pos) <= self.POINT_RADIUS

    @staticmethod
    def _dist(a: tuple[int, int], b: tuple[int, int]) -> float:
        return math.hypot(a[0] - b[0], a[1] - b[1])

    def _new_ab(self) -> tuple[tuple[int, int], tuple[int, int]]:
        """Generate a new random A→B pair within distance constraints.

        Strategy: clamp dist_min/dist_max to the canvas diagonal.  When the
        required distance is large relative to the canvas, bias A towards edges
        and corners so that long-distance B positions become reachable.  The
        fallback randomly picks from the four corner pairs with jitter to ensure
        variety even in the degenerate case.
        """
        margin = self.POINT_RADIUS + 5
        x_lo, x_hi = margin, self.screen_w - margin
        y_lo, y_hi = margin, self.screen_h - margin
        w_inner, h_inner = x_hi - x_lo, y_hi - y_lo
        max_possible = int(math.hypot(w_inner, h_inner))
        eff_max = min(self.dist_max, max_possible)
        eff_min = min(self.dist_min, eff_max)

        # Determine how "tight" the distance requirement is relative to canvas.
        # When ratio > 0.7, purely random A rarely works — bias towards edges.
        tightness = eff_min / max_possible if max_possible > 0 else 1.0

        for _ in range(500):
            if tightness > 0.7:
                # Bias A towards edges/corners: pick from a ring near the border
                side = random.choice(["top", "bottom", "left", "right"])
                edge_band = max(int(w_inner * 0.15), 1)
                if side == "top":
                    ax = random.randint(x_lo, x_hi)
                    ay = random.randint(y_lo, y_lo + edge_band)
                elif side == "bottom":
                    ax = random.randint(x_lo, x_hi)
                    ay = random.randint(y_hi - edge_band, y_hi)
                elif side == "left":
                    ax = random.randint(x_lo, x_lo + edge_band)
                    ay = random.randint(y_lo, y_hi)
                else:
                    ax = random.randint(x_hi - edge_band, x_hi)
                    ay = random.randint(y_lo, y_hi)
            else:
                ax = random.randint(x_lo, x_hi)
                ay = random.randint(y_lo, y_hi)

            # Compute the farthest reachable distance from (ax, ay) within bounds
            reach = max(
                math.hypot(ax - x_lo, ay - y_lo),
                math.hypot(ax - x_hi, ay - y_lo),
                math.hypot(ax - x_lo, ay - y_hi),
                math.hypot(ax - x_hi, ay - y_hi),
            )
            if reach < eff_min:
                continue

            local_max = min(eff_max, int(reach))

            # Try several angles from this A
            for _ in range(30):
                angle = random.uniform(0, 2 * math.pi)
                dist = random.randint(eff_min, local_max)
                bx = int(ax + dist * math.cos(angle))
                by = int(ay + dist * math.sin(angle))
                if x_lo <= bx <= x_hi and y_lo <= by <= y_hi:
                    return (ax, ay), (bx, by)

        # Fallback: pick a random corner pair with jitter for variety
        corners = [(x_lo, y_lo), (x_hi, y_lo), (x_lo, y_hi), (x_hi, y_hi)]
        pairs = [(corners[i], corners[j])
                 for i in range(4) for j in range(i + 1, 4)
                 if self._dist(corners[i], corners[j]) >= eff_min]
        if not pairs:
            # All pairs too short — pick the longest pair
            pairs = [(corners[i], corners[j])
                     for i in range(4) for j in range(i + 1, 4)]
            pairs.sort(key=lambda p: self._dist(p[0], p[1]), reverse=True)
            pairs = pairs[:1]

        ca, cb = random.choice(pairs)
        # Add jitter so it's not identical each time
        jitter = max(margin, int(min(w_inner, h_inner) * 0.08))
        ax = ca[0] + random.randint(-jitter, jitter)
        ay = ca[1] + random.randint(-jitter, jitter)
        bx = cb[0] + random.randint(-jitter, jitter)
        by = cb[1] + random.randint(-jitter, jitter)
        # Clamp back into bounds
        ax = max(x_lo, min(x_hi, ax))
        ay = max(y_lo, min(y_hi, ay))
        bx = max(x_lo, min(x_hi, bx))
        by = max(y_lo, min(y_hi, by))
        return (ax, ay), (bx, by)
