"""hCaptcha tile classifier.

Accepts a tile image + instruction text, returns yes/no classification.
Uses Florence-2 for phrase grounding / visual question answering.

This solver is designed for the POST /classify endpoint which the
Playwright bot calls for each tile in the hCaptcha grid.
"""

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


class HCaptchaSolver(BaseSolver):
    name = "hcaptcha"
    captcha_type = "hcaptcha"

    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()

    def attempts(self):
        return [
            self._florence2_classify,
            self._moondream_classify,
        ]

    def _florence2_classify(self) -> SolveAttempt:
        """Classify tile using Florence-2."""
        if self._img is None:
            return SolveAttempt(
                answer="no",
                confidence=0.0,
                solver_name="hcaptcha.florence2",
                error="no image",
            )
        if not self.ctx.florence._loaded:
            try:
                self.ctx.florence.load()
            except Exception as exc:
                return SolveAttempt(
                    answer="no",
                    confidence=0.0,
                    solver_name="hcaptcha.florence2",
                    error=f"load failed: {exc}",
                )

        try:
            import torch

            # Use Florence-2's caption + phrase grounding to classify
            # First, get a caption of the image
            prompt = "<CAPTION>"
            inputs = self.ctx.florence._processor(
                text=prompt, images=self._img, return_tensors="pt"
            ).to(self.ctx.florence._model.device)

            with torch.no_grad():
                gen = self.ctx.florence._model.generate(
                    input_ids=inputs["input_ids"],
                    pixel_values=inputs["pixel_values"].to(self.ctx.florence._model.dtype),
                    max_new_tokens=64,
                    num_beams=3,
                    do_sample=False,
                )
            caption = self.ctx.florence._processor.batch_decode(
                gen, skip_special_tokens=False
            )[0]
            caption_parsed = self.ctx.florence._processor.post_process_generation(
                caption, task="<CAPTION>", image_size=(self._img.width, self._img.height)
            )
            caption_text = str(caption_parsed.get("<CAPTION>", "")).lower()

            # Now ask if the hint matches the caption
            hint_lower = self._hint.lower()
            if not hint_lower:
                # No hint - return the caption as answer
                return SolveAttempt(
                    answer=caption_text,
                    confidence=0.4,
                    solver_name="hcaptcha.florence2",
                    metadata={"caption": caption_text},
                )

            # Check if the hint words appear in the caption
            hint_words = hint_lower.split()
            matches = sum(1 for w in hint_words if w in caption_text)
            ratio = matches / len(hint_words) if hint_words else 0

            is_match = ratio >= 0.5  # At least half the hint words match
            return SolveAttempt(
                answer="yes" if is_match else "no",
                confidence=0.75 if is_match else 0.65,
                solver_name="hcaptcha.florence2",
                metadata={"caption": caption_text, "match_ratio": ratio},
            )

        except Exception as exc:
            return SolveAttempt(
                answer="no",
                confidence=0.0,
                solver_name="hcaptcha.florence2",
                error=str(exc),
            )

    def _moondream_classify(self) -> SolveAttempt:
        """Classify tile using Moondream2 VQA."""
        if self._img is None:
            return SolveAttempt(
                answer="no",
                confidence=0.0,
                solver_name="hcaptcha.moondream",
                error="no image",
            )

        try:
            hint = self._hint or "the main object"
            question = f"Does this image contain {hint}? Answer yes or no only."
            out = self.ctx.moondream.query(self._img, question, max_tokens=10)

            is_yes = out.strip().lower().startswith("yes")
            return SolveAttempt(
                answer="yes" if is_yes else "no",
                confidence=0.70 if is_yes else 0.60,
                solver_name="hcaptcha.moondream",
                metadata={"raw_answer": out},
            )
        except Exception as exc:
            return SolveAttempt(
                answer="no",
                confidence=0.0,
                solver_name="hcaptcha.moondream",
                error=str(exc),
            )


def classify_tile(image_b64: str, instruction: str, ctx) -> dict:
    """Quick classifier for a single tile. Used by POST /classify.

    Args:
        image_b64: Base64-encoded tile image.
        instruction: hCaptcha instruction (e.g. "Find all items that were made by people").
        ctx: SolveContext with loaded engines.

    Returns:
        dict with "match" (bool), "confidence" (float), "caption" (str).
    """
    solver = HCaptchaSolver(ctx)
    solver.prepare(image_b64, None, instruction)

    # Try Florence-2 first, then Moondream
    for attempt_fn in solver.attempts():
        result = attempt_fn()
        if result.confidence >= 0.5:
            return {
                "match": result.answer.lower() == "yes",
                "confidence": result.confidence,
                "caption": result.metadata.get("caption", result.answer),
                "solver": result.solver_name,
            }

    return {
        "match": False,
        "confidence": 0.0,
        "caption": "",
        "solver": "hcaptcha.none",
    }