"""Moondream2 engine for VQA on captcha images.

Moondream2 is a 1.9B vision-language model that runs on CPU (~2GB RAM).
It's small but good at visual question answering, which makes it
suitable for image-grid captcha style queries like:
  'Which squares contain a traffic light?'
  'Is there a bicycle in this image? answer yes or no'
"""

from __future__ import annotations

import os
from typing import Optional

from captcha_solver.engines.base import BaseEngine
from captcha_solver.config import get_settings


class MoondreamEngine(BaseEngine):
    name = "moondream2"

    def __init__(self) -> None:
        super().__init__()
        self._model = None
        self._tokenizer = None

    def _do_load(self) -> None:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        s = get_settings()
        os.environ.setdefault("HF_HOME", str(s.cache_dir / "hf"))

        self._model = AutoModelForCausalLM.from_pretrained(
            s.moondream_model,
            trust_remote_code=True,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True,
        ).to(s.moondream_device).eval()
        self._tokenizer = AutoTokenizer.from_pretrained(
            s.moondream_model, trust_remote_code=True
        )

    def _do_unload(self) -> None:
        self._model = None
        self._tokenizer = None

    def query(self, pil_image, question: str, max_tokens: int = 80) -> str:
        """Ask a yes/no or short-answer question about an image."""
        if not self._loaded:
            self.load()
        import torch
        assert self._model is not None and self._tokenizer is not None

        with torch.no_grad():
            enc = self._model.encode_image(pil_image)
            if hasattr(self._model, "decode_question"):
                prompt = self._model.decode_question(question)
                out = self._model.generate(
                    image_embeds=enc,
                    prompt=prompt,
                    tokenizer=self._tokenizer,
                    max_new_tokens=max_tokens,
                    do_sample=False,
                )
                return self._tokenizer.decode(out[0], skip_special_tokens=True).strip()
            elif hasattr(self._model, "answer_question"):
                return self._model.answer_question(
                    image_embeds=enc, question=question, tokenizer=self._tokenizer
                ).strip()
            elif hasattr(self._model, "query"):
                return self._model.query(image=enc, question=question).get("answer", "").strip()
            else:
                return self._model.generate(
                    inputs=enc,
                    question=question,
                    tokenizer=self._tokenizer,
                    max_new_tokens=max_tokens,
                )