cv_challenge.py

"""
CV-based coordinate calculator for hCaptcha challenges.

Uses OpenCV, numpy, PIL and scipy instead of relying solely on Gemini's
spatial guesses — gives pixel-accurate drag/click coordinates.

Supported challenges:
  - image_drag_multi  : "drag blocks to cover all icon patterns"
                        "drag the TWO shapes into the correct EMPTY SPACES"
  - image_drag_single : "drag element to most similar"
  - image_label_*     : path intersection / animal-meet (heuristic assist)
"""

import io
import os
from collections import Counter
from pathlib import Path
from typing import List, Optional, Tuple

import cv2
import numpy as np
from PIL import Image, ImageDraw
from scipy.spatial.distance import cosine

# ─── Type aliases ─────────────────────────────────────────────────────────────
PageXY = Tuple[float, float]
DragPath = Tuple[PageXY, PageXY]         # (from_page_xy, to_page_xy)
BBox = dict                               # {'x','y','width','height'}


# ══════════════════════════════════════════════════════════════════════════════
# Basic image utilities
# ══════════════════════════════════════════════════════════════════════════════

def load_img(path: str | Path) -> np.ndarray:
    img = cv2.imread(str(path))
    if img is None:
        raise FileNotFoundError(f"Cannot load image: {path}")
    return img


def bytes_to_cv2(data: bytes) -> np.ndarray:
    arr = np.frombuffer(data, np.uint8)
    return cv2.imdecode(arr, cv2.IMREAD_COLOR)


def cv2_to_pil(img: np.ndarray) -> Image.Image:
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))


def pil_to_bytes(img: Image.Image, fmt: str = "PNG") -> bytes:
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    return buf.getvalue()


def img_to_page(pixel_x: float, pixel_y: float,
                img_w: int, img_h: int, bbox: BBox) -> PageXY:
    """Convert pixel position in screenshot → absolute page coordinate."""
    px = bbox["x"] + (pixel_x / img_w) * bbox["width"]
    py = bbox["y"] + (pixel_y / img_h) * bbox["height"]
    return (px, py)


# ══════════════════════════════════════════════════════════════════════════════
# Feature / similarity helpers
# ══════════════════════════════════════════════════════════════════════════════

def color_histogram(img_bgr: np.ndarray, bins: int = 32) -> np.ndarray:
    """Normalised BGR color histogram vector."""
    hists = [
        cv2.calcHist([img_bgr], [c], None, [bins], [0, 256]).flatten()
        for c in range(3)
    ]
    h = np.concatenate(hists)
    s = h.sum()
    return h / s if s > 0 else h


def hist_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity of color histograms (1 = identical, 0 = opposite)."""
    ha, hb = color_histogram(a), color_histogram(b)
    if ha.sum() == 0 or hb.sum() == 0:
        return 0.0
    return float(1.0 - cosine(ha, hb))


def template_match_score(template: np.ndarray, target: np.ndarray) -> float:
    """
    Normalised cross-correlation between template and target.
    Resizes template to target size for fair comparison.
    Returns best match score (0–1).
    """
    if template.size == 0 or target.size == 0:
        return 0.0
    t = cv2.resize(template, (target.shape[1], target.shape[0]))
    t_gray = cv2.cvtColor(t, cv2.COLOR_BGR2GRAY)
    g_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)
    res = cv2.matchTemplate(g_gray, t_gray, cv2.TM_CCOEFF_NORMED)
    _, max_val, _, _ = cv2.minMaxLoc(res)
    return float(max_val)


def combined_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Weighted combination of histogram + template similarity."""
    return 0.6 * hist_similarity(a, b) + 0.4 * template_match_score(a, b)


# ══════════════════════════════════════════════════════════════════════════════
# Contour / region detection
# ══════════════════════════════════════════════════════════════════════════════

def detect_significant_regions(img: np.ndarray,
                                min_area: int = 400,
                                max_area_ratio: float = 0.4
                                ) -> List[Tuple[int, int, int, int]]:
    """
    Find bounding boxes of visually significant contours.
    Returns sorted list of (x, y, w, h).
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    img_area = img.shape[0] * img.shape[1]
    boxes = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        area = w * h
        if min_area < area < max_area_ratio * img_area:
            boxes.append((x, y, w, h))
    return sorted(boxes, key=lambda b: (b[1], b[0]))  # top-left to bottom-right


def split_into_grid(img: np.ndarray,
                    cols: int, rows: int
                    ) -> List[Tuple[Tuple[int,int,int,int], np.ndarray]]:
    """
    Divide image into a regular grid.
    Returns list of ((x,y,w,h), cell_image) pairs.
    """
    h, w = img.shape[:2]
    cw, ch = w // cols, h // rows
    cells = []
    for r in range(rows):
        for c in range(cols):
            x, y = c * cw, r * ch
            cell = img[y:y+ch, x:x+cw]
            cells.append(((x, y, cw, ch), cell))
    return cells


def cell_variance(cell: np.ndarray) -> float:
    """Pixel variance — low variance means the cell looks empty/uniform."""
    return float(np.var(cv2.cvtColor(cell, cv2.COLOR_BGR2GRAY)))


# ══════════════════════════════════════════════════════════════════════════════
# Challenge: "drag blocks to cover all icon patterns"
# ══════════════════════════════════════════════════════════════════════════════

def _find_panel_split(img: np.ndarray) -> int:
    """
    Heuristically find the vertical split between the puzzle grid (left)
    and the draggable-block panel (right).
    Uses a vertical line of high contrast as the separator.
    Returns pixel column index of split.
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Look for a dark/bright vertical separator in the right 40%
    h, w = gray.shape
    search = gray[:, int(w * 0.55):int(w * 0.85)]
    col_std = np.std(search, axis=0)
    # The separator column often has very low std (solid color bar)
    sep_local = int(np.argmin(col_std))
    return int(w * 0.55) + sep_local


def _find_colored_blocks(img: np.ndarray, w: int, h: int,
                          bbox: BBox) -> List[Tuple[PageXY, np.ndarray]]:
    """
    Detect draggable colored blocks across the whole image using HSV ranges.
    Handles vertical stacks (purple panel) and horizontal rows (green panel).
    Returns list of (page_pos, crop_img).
    """
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    blocks = []

    # Build a combined saturation mask for all draggable-block colors
    # (purple H≈120-160, green/teal H≈40-110) and filter out backgrounds.
    combined_lo = np.array([40,  50,  50])
    combined_hi = np.array([165, 255, 255])
    mask = cv2.inRange(hsv, combined_lo, combined_hi)

    # Exclude header (top ≈ 22% of image, usually the teal instruction bar)
    header_px = max(100, h // 5)
    mask[:header_px, :] = 0

    # Exclude grid area (left ≈ 55% of image where icons live)
    grid_split = int(w * 0.55)
    mask[:, :grid_split] = 0

    cont, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Sort all contours by area, largest first.
    # Filter: min area 3000 to skip noise; max area 50000 to skip background.
    panels = sorted(
        [(cv2.boundingRect(c), cv2.contourArea(c)) for c in cont
         if 3000 < cv2.contourArea(c) < 50000],
        key=lambda t: t[1], reverse=True,
    )

    RATIO_THRESH = 1.25   # aspect-ratio threshold to call a panel "non-square"

    def _add_panel(x: int, y: int, pw: int, ph: int) -> None:
        aspect = ph / pw if pw > 0 else 1.0
        if pw > ph * RATIO_THRESH:                  # horizontal → split cols
            n = max(1, round(aspect ** -1))         # n = round(pw/ph)
            n = max(1, round(pw / ph))
            sw = pw // n
            for i in range(n):
                cx = x + i * sw + sw // 2
                cy = y + ph // 2
                crop = img[y:y+ph, x+i*sw:x+(i+1)*sw]
                blocks.append((img_to_page(cx, cy, w, h, bbox), crop))
        elif ph > pw * RATIO_THRESH:                # vertical → split rows
            n = max(1, round(ph / pw))
            sh = ph // n
            for i in range(n):
                cx = x + pw // 2
                cy = y + i * sh + sh // 2
                crop = img[y+i*sh:y+(i+1)*sh, x:x+pw]
                blocks.append((img_to_page(cx, cy, w, h, bbox), crop))
        else:                                        # square-ish → single block
            cx, cy = x + pw // 2, y + ph // 2
            crop = img[y:y+ph, x:x+pw]
            blocks.append((img_to_page(cx, cy, w, h, bbox), crop))

    for (x, y, pw, ph), area in panels[:6]:
        # SKIP square-ish blobs that are very large: these are usually combined
        # overlap regions of two panels touching each other.  Genuine single
        # blocks are small; genuine panels have a clear aspect ratio.
        aspect_ratio = max(pw / ph, ph / pw) if min(pw, ph) > 0 else 1.0
        max_single_block_area = 12_000
        if aspect_ratio < RATIO_THRESH and area > max_single_block_area:
            # Too big to be a single block, not elongated enough to be a panel
            # → skip (handled indirectly by the panels with clear ratios)
            continue
        _add_panel(x, y, pw, ph)

    # De-duplicate: remove positions within 30 px of each other
    unique: List[Tuple[PageXY, np.ndarray]] = []
    for pos, crop in blocks:
        ix = int((pos[0] - bbox["x"]) / bbox["width"] * w)
        iy = int((pos[1] - bbox["y"]) / bbox["height"] * h)
        dup = any(
            abs(ix - int((p[0] - bbox["x"]) / bbox["width"] * w)) < 30 and
            abs(iy - int((p[1] - bbox["y"]) / bbox["height"] * h)) < 30
            for p, _ in unique
        )
        if not dup:
            unique.append((pos, crop))

    blocks.clear()
    blocks.extend(unique)
    return blocks


def _find_icon_positions(img: np.ndarray, w: int, h: int,
                          bbox: BBox,
                          thresh: int = 200,
                          grid_x_max_frac: float = 0.60,
                          header_skip: int = 100,
                          min_area: int = 300,
                          max_area: int = 8000
                          ) -> List[Tuple[PageXY, np.ndarray, int]]:
    """
    Find white icon patterns in the grid area using brightness threshold.
    Returns list of (page_pos, crop, area) sorted by area descending.
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, bright = cv2.threshold(gray, thresh, 255, cv2.THRESH_BINARY)

    # Restrict to grid region: left fraction of image, below header
    roi_mask = np.zeros_like(bright)
    roi_mask[header_skip:h, :int(w * grid_x_max_frac)] = 255
    bright = cv2.bitwise_and(bright, roi_mask)

    cont, _ = cv2.findContours(bright, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    icons = []
    for cnt in cont:
        x, y, cw, ch = cv2.boundingRect(cnt)
        area = cw * ch
        if min_area < area < max_area:
            cx, cy = x + cw // 2, y + ch // 2
            crop = img[y:y+ch, x:x+cw]
            icons.append((img_to_page(cx, cy, w, h, bbox), crop, area))

    icons.sort(key=lambda ic: ic[2], reverse=True)
    return icons


def solve_drag_cover(screenshot_path: str | Path,
                     bbox: BBox,
                     n_blocks: int = 2
                     ) -> Optional[List[DragPath]]:
    """
    Solver for: "drag blocks to cover all icon patterns"
                "drag the TWO shapes into the correct EMPTY SPACES"

    Algorithm
    ---------
    1. Find draggable colored blocks (purple/green) anywhere in the image
       using HSV color detection.  Panels are split into individual sub-blocks.
    2. Find white icon positions in the LEFT grid region using brightness
       threshold — these are the exact pixels of each icon to be covered.
    3. Match each block to the most-similar icon via combined_similarity().
    4. Return drag paths (FROM block → TO icon).
    """
    img = load_img(screenshot_path)
    h, w = img.shape[:2]

    # ── 1. Find draggable blocks by color ────────────────────────────────
    blocks = _find_colored_blocks(img, w, h, bbox)
    if not blocks:
        print("[cv2] no colored blocks found")
        return None

    # ── 2. Find icon positions by brightness ─────────────────────────────
    icons = _find_icon_positions(img, w, h, bbox)
    if not icons:
        print("[cv2] no icons found in grid area")
        return None

    n = min(len(blocks), len(icons), max(n_blocks, len(blocks)))
    print(f"[cv2] blocks={len(blocks)} icons={len(icons)} dragging={n}")

    # ── 3. Match blocks to icons by similarity ────────────────────────────
    used_icons: set = set()
    paths: List[DragPath] = []

    for blk_page, blk_crop in blocks[:n]:
        best_score, best_idx = -1.0, 0
        for i, (ico_page, ico_crop, _) in enumerate(icons[:n + 2]):
            if i in used_icons:
                continue
            score = combined_similarity(blk_crop, ico_crop)
            if score > best_score:
                best_score, best_idx = score, i
        used_icons.add(best_idx)
        paths.append((blk_page, icons[best_idx][0]))

    return paths if paths else None


# ══════════════════════════════════════════════════════════════════════════════
# Challenge: "drag element on the left to the one that is most similar"
# ══════════════════════════════════════════════════════════════════════════════

def solve_similarity_drag(screenshot_path: str | Path,
                          bbox: BBox
                          ) -> Optional[DragPath]:
    """
    Solver for: "Please drag the element on the left to the one that is most similar"

    Algorithm
    ---------
    1. Find the reference element (left panel, tagged with '+Move').
    2. Find all candidate elements (right panel, 2×2 or 3×1 grid).
    3. Use combined_similarity() to score each candidate.
    4. Return drag path from reference to best candidate.
    """
    img = load_img(screenshot_path)
    h, w = img.shape[:2]

    split = int(w * 0.30)          # reference occupies left ~30%
    ref_panel = img[:, :split]
    cand_panel = img[:, split:]
    cp_w = cand_panel.shape[1]

    # Reference: largest contour in left panel
    ref_regions = detect_significant_regions(ref_panel, min_area=500)
    if not ref_regions:
        ref_regions = [(0, 0, split, h)]
    rx, ry, rw2, rh2 = sorted(ref_regions, key=lambda b: b[2]*b[3], reverse=True)[0]
    ref_crop = ref_panel[ry:ry+rh2, rx:rx+rw2]
    ref_page = img_to_page(rx + rw2//2, ry + rh2//2, w, h, bbox)

    # Candidates: try 2×2 grid first, fall back to detected regions
    cand_cells = split_into_grid(cand_panel, cols=2, rows=2)

    best_score, best_page = -1.0, (split + cp_w//2, h//2)
    for (cx, cy, cw2, ch2), cand_crop in cand_cells:
        score = combined_similarity(ref_crop, cand_crop)
        if score > best_score:
            best_score = score
            abs_cx = split + cx + cw2 // 2
            abs_cy = cy + ch2 // 2
            best_page = img_to_page(abs_cx, abs_cy, w, h, bbox)

    return ref_page, best_page


# ══════════════════════════════════════════════════════════════════════════════
# Challenge: "Click the circle where the animals will meet"
# ══════════════════════════════════════════════════════════════════════════════

def solve_path_intersection(screenshot_path: str | Path,
                            bbox: BBox
                            ) -> Optional[PageXY]:
    """
    Solver for: "Click the circle where the animals will meet"

    Heuristic: find the intersection region by looking for the area where
    two distinct colored path lines cross. Uses Hough line segments.

    Returns click page coordinate.
    """
    img = load_img(screenshot_path)
    h, w = img.shape[:2]

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=40,
                             minLineLength=30, maxLineGap=15)
    if lines is None:
        # No lines found — click center of image as fallback
        return img_to_page(w//2, h//2, w, h, bbox)

    # Collect endpoints of all line segments
    points = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        points.extend([(x1, y1), (x2, y2)])

    # Find the densest cluster of endpoints using a heat map
    heat = np.zeros((h, w), dtype=np.float32)
    for px, py in points:
        if 0 <= px < w and 0 <= py < h:
            cv2.circle(heat, (px, py), radius=20, color=1.0, thickness=-1)

    # Blur to smooth, find peak
    heat = cv2.GaussianBlur(heat, (31, 31), 0)
    _, _, _, max_loc = cv2.minMaxLoc(heat)
    cx, cy = max_loc

    return img_to_page(cx, cy, w, h, bbox)


# ══════════════════════════════════════════════════════════════════════════════
# Dispatcher: detect prompt → run correct CV solver
# ══════════════════════════════════════════════════════════════════════════════

PROMPT_HANDLERS = {
    # drag blocks / shapes to empty spaces
    "empty spaces":     "drag_cover",
    "cover all":        "drag_cover",
    "cover the":        "drag_cover",
    "drag the blocks":  "drag_cover",
    "drag blocks":      "drag_cover",
    # drag to similar
    "most similar":     "similarity",
    "similar":          "similarity",
    # path intersection
    "animals will meet":"path",
    "circle where":     "path",
    "will meet":        "path",
}

# Direct mapping from ChallengeTypeEnum values to CV handler
TYPE_HANDLERS = {
    "image_drag_multi":  "drag_cover",
    "image_drag_single": "similarity",
    # label types handled well by Gemini — no CV override
    "image_label_single_select": None,
    "image_label_multi_select":  None,
}


def dispatch_by_type(job_type_value: str,
                     screenshot_path: str | Path,
                     bbox: BBox,
                     n_blocks: int = 2
                     ) -> Optional[object]:
    """
    Dispatch CV solver using the ChallengeTypeEnum string value directly.
    More reliable than prompt keyword matching.

    Returns:
      drag_cover  → List[DragPath]  [(from_xy, to_xy), ...]
      similarity  → DragPath        (from_xy, to_xy)
      None        → let Gemini handle it
    """
    handler = TYPE_HANDLERS.get(job_type_value)
    if handler is None:
        return None

    try:
        if handler == "drag_cover":
            result = solve_drag_cover(screenshot_path, bbox, n_blocks=n_blocks)
            print(f"[cv2] drag_cover result: {result}")
            return result
        elif handler == "similarity":
            result = solve_similarity_drag(screenshot_path, bbox)
            print(f"[cv2] similarity result: {result}")
            return result
    except Exception as e:
        print(f"[cv_challenge] dispatch_by_type error ({handler}): {e}")
        return None


def dispatch(prompt: str,
             screenshot_path: str | Path,
             bbox: BBox,
             n_blocks: int = 2
             ) -> Optional[object]:
    """
    Dispatch by prompt keyword (legacy — prefer dispatch_by_type).
    """
    p = prompt.lower()
    handler = None
    for kw, kind in PROMPT_HANDLERS.items():
        if kw in p:
            handler = kind
            break

    if handler is None:
        return None

    try:
        if handler == "drag_cover":
            return solve_drag_cover(screenshot_path, bbox, n_blocks=n_blocks)
        elif handler == "similarity":
            return solve_similarity_drag(screenshot_path, bbox)
        elif handler == "path":
            return solve_path_intersection(screenshot_path, bbox)
    except Exception as e:
        print(f"[cv_challenge] dispatch error ({handler}): {e}")
        return None


# ══════════════════════════════════════════════════════════════════════════════
# Debug helper: annotate screenshot with computed coordinates
# ══════════════════════════════════════════════════════════════════════════════

def annotate(screenshot_path: str | Path,
             paths: List[DragPath],
             bbox: BBox,
             out_path: str | Path = "/tmp/cv_debug.png"):
    """Draw computed drag arrows on screenshot for visual debugging."""
    img = load_img(screenshot_path)
    h, w = img.shape[:2]
    pil = cv2_to_pil(img)
    draw = ImageDraw.Draw(pil)

    def to_img(px, py):
        ix = (px - bbox["x"]) / bbox["width"] * w
        iy = (py - bbox["y"]) / bbox["height"] * h
        return int(ix), int(iy)

    for (fx, fy), (tx, ty) in paths:
        fi = to_img(fx, fy)
        ti = to_img(tx, ty)
        draw.ellipse([fi[0]-8, fi[1]-8, fi[0]+8, fi[1]+8], fill="lime")
        draw.ellipse([ti[0]-8, ti[1]-8, ti[0]+8, ti[1]+8], fill="red")
        draw.line([fi, ti], fill="yellow", width=3)

    pil.save(str(out_path))
    print(f"[cv_challenge] debug image → {out_path}")