"""Real (tiny) model backends, small enough to run on a laptop CPU.

Two of them, because they trade off differently:

* :class:`TextPlannerBrain` - a ~0.5B *text* model over the OCR index. No vision,
  no big download, and the element index does all the grounding, which is why it
  works at this size at all.
* :class:`VLMBrain` - a tiny *vision* model that sees the screenshot itself
  (256M-class models can read a button label from pixels, but reason poorly, so the
  menu still carries most of the decision).

Both inherit menu building, prompt shape and lenient parsing from
:class:`~tiny_computer_agent.brain.ModelBrain`, so swapping weights never changes the
harness contract.
"""

from __future__ import annotations

import os
from typing import Any

from tiny_computer_agent.brain import BrainError, ModelBrain
from tiny_computer_agent.decoding import MenuMask
from tiny_computer_agent.observe import Observation

DEFAULT_TEXT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
DEFAULT_VISION_MODEL = "HuggingFaceTB/SmolVLM-256M-Instruct"


def _pair_margin(logits: Any, yes_ids: Any, no_ids: Any) -> float:
    """log P(yes) - log P(no) for one position, over the whole vocabulary."""
    import torch

    yes = torch.logsumexp(logits[list(yes_ids)], dim=0)
    no = torch.logsumexp(logits[list(no_ids)], dim=0)
    return float((yes - no).item())


def _masked_generation(model: Any, allowed_ids: Any, eos_token_id: Any) -> Any:
    """A transformers ``LogitsProcessorList`` holding the menu mask.

    Imported here, not at module scope, so the harness still imports with no
    transformers installed at all.
    """
    try:
        from transformers import LogitsProcessorList
    except ImportError:                                   # pragma: no cover - env
        raise BrainError(
            "constrained decoding needs transformers: pip install 'tiny-computer-agent[vision]'"
        )
    return LogitsProcessorList([MenuMask(allowed_ids, eos_token_id=eos_token_id)])


def resolve_device(prefer: str | None = None) -> str:
    """Pick a torch device: explicit preference, then CUDA, then CPU."""
    try:
        import torch
    except ImportError:  # pragma: no cover - depends on the machine
        raise BrainError("model backends need torch: pip install 'tiny-computer-agent[vision]'")
    if prefer:
        return prefer
    return "cuda" if torch.cuda.is_available() else "cpu"


def pick_dtype(device: str) -> Any:
    import torch

    if device.startswith("cuda"):
        return torch.bfloat16
    return torch.float32


def model_dependencies() -> tuple[bool, str]:
    """(available, detail) for preflight and the doctor command."""
    try:
        import transformers
    except ImportError:
        return False, "transformers is not installed (pip install 'tiny-computer-agent[vision]')"
    try:
        import torch
    except ImportError:
        return False, "torch is not installed (pip install 'tiny-computer-agent[vision]')"
    return True, f"transformers {transformers.__version__}, torch {torch.__version__}"


def model_cached(model_id: str) -> bool:
    """True when the weights are already in the local Hugging Face cache."""
    try:
        from huggingface_hub import constants
    except ImportError:  # pragma: no cover
        return False
    folder = os.path.join(constants.HF_HUB_CACHE, "models--" + model_id.replace("/", "--"))
    return os.path.isdir(folder)


class TextPlannerBrain(ModelBrain):
    """Causal-LM planner over the OCR index: no pixels in, one menu number out."""

    name = "text-planner"
    requires_index = True
    supports_constrained = True

    def __init__(self, model_id: str = DEFAULT_TEXT_MODEL, *, device: str | None = None,
                 **kwargs: Any) -> None:
        super().__init__(model_id, **kwargs)
        self.device = device
        self._tokenizer: Any = None
        self._model: Any = None

    def _load(self) -> None:
        if self._model is not None:
            return
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except ImportError as exc:  # pragma: no cover - depends on the machine
            raise BrainError(
                "text planner needs transformers and torch: "
                "pip install 'tiny-computer-agent[vision]'"
            ) from exc
        self.device = resolve_device(self.device)
        self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        self._model = AutoModelForCausalLM.from_pretrained(
            self.model_id, dtype=pick_dtype(self.device)
        ).to(self.device).eval()

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

    def _tokenizer_for_choice(self) -> Any:
        self._load()
        return self._tokenizer

    def _generate_constrained(self, prompt: str, observation: Observation,
                              allowed_ids: Any, max_new_tokens: int,
                              *, sample: bool = False) -> str:
        """One short generation that can only contain a menu number."""
        import torch

        self._load()
        messages = [{"role": "user", "content": prompt}]
        text = self._tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        inputs = self._tokenizer(text, return_tensors="pt").to(self.device)
        settings: dict[str, Any] = {
            "max_new_tokens": max_new_tokens,
            "pad_token_id": self._tokenizer.eos_token_id,
            "logits_processor": _masked_generation(
                self._model, allowed_ids, self._tokenizer.eos_token_id),
        }
        if sample:
            settings.update(do_sample=True, temperature=0.7, top_p=0.9)
        else:
            settings["do_sample"] = False
        with torch.inference_mode():
            output = self._model.generate(**inputs, **settings)
        return self._tokenizer.decode(
            output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
        ).strip()

    def _score_pair(self, prompt: str, observation: Observation,
                    yes_ids: Any, no_ids: Any) -> float:
        import torch

        self._load()
        messages = [{"role": "user", "content": prompt}]
        text = self._tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False)
        inputs = self._tokenizer(text, return_tensors="pt").to(self.device)
        with torch.inference_mode():
            out = self._model.generate(
                **inputs, max_new_tokens=1, do_sample=False, return_dict_in_generate=True,
                output_scores=True, pad_token_id=self._tokenizer.eos_token_id)
        return _pair_margin(out.scores[0][0], yes_ids, no_ids)

    def _generate(self, prompt: str, observation: Observation) -> str:
        import torch

        self._load()
        messages = [{"role": "user", "content": prompt}]
        text = self._tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        inputs = self._tokenizer(text, return_tensors="pt").to(self.device)
        with torch.inference_mode():
            output = self._model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                do_sample=False,
                pad_token_id=self._tokenizer.eos_token_id,
            )
        return self._tokenizer.decode(
            output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
        ).strip()


class VLMBrain(ModelBrain):
    """Tiny vision-language model: sees the screenshot, answers with a menu number."""

    name = "vlm"
    requires_index = False
    supports_constrained = True

    def __init__(self, model_id: str = DEFAULT_VISION_MODEL, *, device: str | None = None,
                 **kwargs: Any) -> None:
        super().__init__(model_id, **kwargs)
        self.device = device
        self._processor: Any = None
        self._model: Any = None

    def _load(self) -> None:
        if self._model is not None:
            return
        try:
            from transformers import AutoModelForImageTextToText, AutoProcessor
        except ImportError as exc:  # pragma: no cover - depends on the machine
            raise BrainError(
                "the vision backend needs transformers and torch: "
                "pip install 'tiny-computer-agent[vision]'"
            ) from exc
        self.device = resolve_device(self.device)
        self._processor = AutoProcessor.from_pretrained(self.model_id)
        self._model = AutoModelForImageTextToText.from_pretrained(
            self.model_id, dtype=pick_dtype(self.device)
        ).to(self.device).eval()

    def unload(self) -> None:
        self._model = None
        self._processor = None

    def _tokenizer_for_choice(self) -> Any:
        self._load()
        return self._processor.tokenizer

    def _generate_constrained(self, prompt: str, observation: Observation,
                              allowed_ids: Any, max_new_tokens: int,
                              *, sample: bool = False) -> str:
        """A vision generation that can only contain a menu number.

        This is the fix for the echo: the model is no longer able to reply with
        the prompt, because the prompt's tokens are masked out of the logits.
        """
        import torch

        image = observation.frame
        if image is None or isinstance(image, str):
            raise BrainError(
                "the vision backend needs a real screenshot but the observation carries "
                f"{type(image).__name__} - run with --live, or use the text planner"
            )
        self._load()
        messages = [{"role": "user", "content": [{"type": "image"},
                                                 {"type": "text", "text": prompt}]}]
        text = self._processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device)
        settings: dict[str, Any] = {
            "max_new_tokens": max_new_tokens,
            "logits_processor": _masked_generation(
                self._model, allowed_ids,
                getattr(self._processor.tokenizer, "eos_token_id", None)),
        }
        if sample:
            settings.update(do_sample=True, temperature=0.7, top_p=0.9)
        else:
            settings["do_sample"] = False
        with torch.inference_mode():
            output = self._model.generate(**inputs, **settings)
        return self._processor.batch_decode(
            output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True
        )[0].strip()

    def _score_pair(self, prompt: str, observation: Observation,
                    yes_ids: Any, no_ids: Any) -> float:
        """log P(yes) - log P(no) for one image question, read off the logits."""
        import torch

        image = observation.frame
        if image is None or isinstance(image, str):
            raise BrainError(
                "the vision backend needs a real screenshot but the observation carries "
                f"{type(image).__name__} - run with --live, or use the text planner"
            )
        self._load()
        messages = [{"role": "user", "content": [{"type": "image"},
                                                 {"type": "text", "text": prompt}]}]
        text = self._processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device)
        with torch.inference_mode():
            out = self._model.generate(**inputs, max_new_tokens=1, do_sample=False,
                                       return_dict_in_generate=True, output_scores=True)
        return _pair_margin(out.scores[0][0], yes_ids, no_ids)

    def _generate(self, prompt: str, observation: Observation) -> str:
        import torch

        image = observation.frame
        if image is None or isinstance(image, str):
            raise BrainError(
                "the vision backend needs a real screenshot but the observation carries "
                f"{type(image).__name__} - run with --live, or use the text planner"
            )
        self._load()
        messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
        text = self._processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = self._processor(text=text, images=[image], return_tensors="pt").to(self.device)
        with torch.inference_mode():
            output = self._model.generate(
                **inputs, max_new_tokens=self.max_new_tokens, do_sample=False
            )
        return self._processor.batch_decode(
            output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True
        )[0].strip()


MODEL_BRAINS = {"text-planner": TextPlannerBrain, "vlm": VLMBrain}


def make_model_brain(name: str = "text-planner", **kwargs: Any) -> ModelBrain:
    """Build a model brain by name, with an actionable error for unknown names."""
    if name not in MODEL_BRAINS:
        raise BrainError(f"unknown brain {name!r}; expected one of {tuple(MODEL_BRAINS)}")
    return MODEL_BRAINS[name](**kwargs)