"""Image-grid captcha solver (e.g. reCAPTCHA v2 'select all squares with X').

We split the image into a 3x3 grid and ask Moondream2 (or Ollama vision)
per-tile 'does this tile contain <hint>?'. Return the indices of tiles
that said yes.

Strategy:
  1. Try ollama vision if enabled (Qwen2-VL 7B best)
  2. Fall back to Moondream2 (CPU, smaller, less accurate)
  3. If neither works, return empty with an error

The hint comes from the page context. If the user provides
hint='traffic lights', we ask per-tile. Common reCAPTCHA categories:
  traffic lights, buses, bicycles, crosswalks, fire hydrants, motorcycles,
  palm trees, stairs, tractors, chimneys
"""

from __future__ import annotations

import re
from typing import Optional

from captcha_solver.solvers.base import BaseSolver, SolveAttempt
from captcha_solver.utils.image import decode_base64_image, image_to_pil


COMMON_CATEGORIES = (
    "traffic light", "bus", "bicycle", "crosswalk", "fire hydrant",
    "motorcycle", "palm tree", "stairs", "tractor", "chimney",
    "boat", "mountain", "car", "truck", "bridge", "house",
)


class ImageGridSolver(BaseSolver):
    name = "image_grid"
    captcha_type = "image_grid"

    def __init__(self, ctx) -> None:
        super().__init__(ctx)
        self._img = None
        self._hint: str = ""

    def prepare(self, image_b64: Optional[str], audio_b64: Optional[str], hint: Optional[str]) -> None:
        if not image_b64:
            self._img = None
            return
        try:
            data = decode_base64_image(image_b64)
            self._img = image_to_pil(data)
        except Exception as exc:
            self._img = None
            self._last_error = f"decode: {exc}"
            return
        self._hint = (hint or "").strip() or self._guess_hint()

    def _guess_hint(self) -> str:
        """If no hint given, return ''; the solver will ask VLM to describe."""
        return ""

    def attempts(self):
        return [
            self._ollama_vision,
            self._moondream,
        ]

    def _split_grid(self, img) -> list:
        w, h = img.size
        tiles: list = []
        for r in range(3):
            for c in range(3):
                box = (c * w // 3, r * h // 3, (c + 1) * w // 3, (r + 1) * h // 3)
                tiles.append(img.crop(box))
        return tiles

    def _format_yes_no(self, raw: str) -> Optional[bool]:
        s = raw.strip().lower()
        if s.startswith("yes") or s.startswith("y") or "yes" in s.split()[:3]:
            return True
        if s.startswith("no") or s.startswith("n") or "no" in s.split()[:3]:
            return False
        return None

    def _ollama_vision(self) -> SolveAttempt:
        if not self.ctx.ollama.enabled or self._img is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="image_grid.ollama", error="ollama disabled")
        tiles = self._split_grid(self._img)
        if self._hint:
            prompt = (
                f"Does this image tile contain a {self._hint}? "
                f"Reply ONLY with 'yes' or 'no'."
            )
        else:
            prompt = (
                "Look at the whole image. It is a 3x3 grid captcha. "
                "Reply with the categories of objects you see in each tile, e.g. '1:car 2:tree 3:car ...'."
            )
        try:
            if self._hint:
                results: list[bool] = []
                for tile in tiles:
                    out = self.ctx.ollama.describe_image(tile, prompt)
                    yn = self._format_yes_no(out)
                    if yn is not None:
                        results.append(yn)
                    else:
                        results.append(False)
                matches = [i + 1 for i, r in enumerate(results) if r]
                return SolveAttempt(
                    answer=",".join(str(m) for m in matches) if matches else "",
                    confidence=0.7 if matches else 0.0,
                    solver_name="image_grid.ollama",
                )
            else:
                out = self.ctx.ollama.describe_image(self._img, prompt)
                return SolveAttempt(answer=out, confidence=0.5, solver_name="image_grid.ollama")
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="image_grid.ollama", error=str(exc))

    def _moondream(self) -> SolveAttempt:
        if self._img is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="image_grid.moondream", error="no image")

        if not self._hint:
            try:
                out = self.ctx.moondream.query(
                    self._img, "Look at this captcha challenge image. Read all the text visible in the image. What does it say?"
                )
            except Exception as exc:
                return SolveAttempt(answer="", confidence=0.0, solver_name="image_grid.moondream", error=str(exc))
            return SolveAttempt(answer=out, confidence=0.4, solver_name="image_grid.moondream")

        # Full-image VQA: ask Moondream to analyze the whole image
        try:
            prompt = (
                f"This is a hCaptcha challenge. The instruction is: '{self._hint}'. "
                f"Look at the image carefully. It shows several tiles each containing a letter and a number. "
                f"Tell me: what specific letter+number combinations should be clicked? "
                f"For example, if the instruction is 'click all the number 3 tiles', "
                f"list the tiles that have the number 3 under them. "
                f"Return the answer as a list of letter+number pairs, e.g. 'A3, B3, C3'."
            )
            out = self.ctx.moondream.query(self._img, prompt)
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="image_grid.moondream", error=str(exc))
        return SolveAttempt(answer=out, confidence=0.5, solver_name="image_grid.moondream")


def infer_hint_from_text(text: str) -> Optional[str]:
    if not text:
        return None
    s = text.lower()
    for cat in COMMON_CATEGORIES:
        if cat in s:
            return cat
    return None