"""End-to-end test for the TRL training pipeline with multimodal VLMs.

Simulates the FULL pipeline on CPU with a tiny model:
  1. Load tiny VLM (CPU, ~100 params)
  2. Wrap in VLMModelWrapper
  3. Build rollout_func via make_waa_rollout_func
  4. Run rollout (generation with screenshot)
  5. Verify model saw pixel_values during generation
  6. Simulate TRL's training forward pass (input_ids only)
  7. Verify wrapper injects cached pixel_values
  8. Verify output is parseable (not garbage)

This test would have caught every bug from the March 29 session:
- Wrong prompt → garbage output → test_output_is_parseable_dsl fails
- Missing pixel_values → blind model → test_forward_gets_pixel_values fails
- Thinking mode → <think> in output → test_no_thinking_tokens fails
- Batch sizing → TRL error → test_rollout_returns_correct_shape fails

Requires torch (CPU only, no GPU). Skipped in CI via @pytest.mark.heavy.
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch
import io

import pytest


@pytest.mark.heavy
class TestTRLE2E:
    """Full pipeline e2e test on CPU with a tiny model."""

    @staticmethod
    def _make_tiny_vlm_and_processor():
        """Build a tiny VLM + processor that work end-to-end.

        The model is ~100 params, runs on CPU in milliseconds.
        The processor mimics Qwen's interface: apply_chat_template + __call__.
        """
        torch = pytest.importorskip("torch")
        import torch.nn as nn

        vocab_size = 200

        class TinyVLM(nn.Module):
            """Minimal VLM that uses pixel_values in forward."""

            def __init__(self):
                super().__init__()
                self.embed = nn.Embedding(vocab_size, 16)
                self.vision_proj = nn.Linear(3, 16)
                self.head = nn.Linear(16, vocab_size)
                self._saw_pixel_values = False

            def forward(self, input_ids, attention_mask=None,
                        pixel_values=None, image_grid_thw=None, **kwargs):
                h = self.embed(input_ids)
                if pixel_values is not None:
                    self._saw_pixel_values = True
                    # Add vision signal — changes logits when image present
                    vis = self.vision_proj(
                        pixel_values.float().mean(dim=(-2, -1))
                    )
                    if vis.dim() == 2:
                        vis = vis.unsqueeze(1)
                    h = h + vis[:, :h.shape[1], :]

                class Out:
                    pass
                out = Out()
                out.logits = self.head(h)
                return out

            def generate(self, input_ids=None, attention_mask=None,
                         pixel_values=None, image_grid_thw=None,
                         max_new_tokens=10, do_sample=True,
                         temperature=1.0, return_dict_in_generate=False,
                         output_scores=False, **kwargs):
                """Minimal generate: run forward, sample greedily."""
                if pixel_values is not None:
                    self._saw_pixel_values = True

                all_ids = input_ids.clone()
                for _ in range(min(max_new_tokens, 20)):
                    out = self(all_ids, pixel_values=pixel_values)
                    next_id = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
                    all_ids = torch.cat([all_ids, next_id], dim=1)

                if return_dict_in_generate:
                    class GenOut:
                        pass
                    result = GenOut()
                    result.sequences = all_ids
                    # Fake scores for logprob computation
                    result.scores = [
                        self.head(self.embed(all_ids[:, i:i+1])).squeeze(1)
                        for i in range(input_ids.shape[1], all_ids.shape[1])
                    ]
                    return result
                return all_ids

        class TinyProcessor:
            """Minimal processor mimicking Qwen's interface."""

            def __init__(self):
                self.tokenizer = self
                # No <think> in template
                self.chat_template = (
                    "{% for msg in messages %}"
                    "<|im_start|>{{ msg.role }}\n{{ msg.content }}<|im_end|>\n"
                    "{% endfor %}"
                    "<|im_start|>assistant\n"
                )

            def apply_chat_template(self, messages, tokenize=False,
                                    add_generation_prompt=True,
                                    enable_thinking=False, **kwargs):
                """Render messages to text."""
                parts = []
                for msg in messages:
                    role = msg["role"]
                    content = msg["content"]
                    if isinstance(content, list):
                        text_parts = [
                            c.get("text", "[image]")
                            for c in content
                            if c.get("type") in ("text", "image")
                        ]
                        content = " ".join(text_parts)
                    parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
                parts.append("<|im_start|>assistant\n")
                return "\n".join(parts)

            def __call__(self, text=None, images=None, return_tensors=None,
                         padding=False):
                """Tokenize text + encode image."""
                # Simple tokenization: each char = 1 token (capped at vocab)
                t = text[0] if isinstance(text, list) else text
                ids = [min(ord(c), 199) for c in t[:50]]  # cap length
                result = {
                    "input_ids": torch.tensor([ids]),
                    "attention_mask": torch.ones(1, len(ids), dtype=torch.long),
                }
                if images:
                    # Create a real pixel_values tensor from the image
                    from PIL import Image
                    img = images[0] if isinstance(images, list) else images
                    if isinstance(img, Image.Image):
                        import numpy as np
                        arr = np.array(img.resize((10, 10)))
                        result["pixel_values"] = torch.tensor(
                            arr, dtype=torch.float32
                        ).permute(2, 0, 1).unsqueeze(0)
                        result["image_grid_thw"] = torch.tensor([[1, 10, 10]])
                return MagicMock(**{k: v for k, v in result.items()},
                                 **{"to": lambda self, d: self,
                                    "get": result.get,
                                    "__contains__": result.__contains__,
                                    "__getitem__": result.__getitem__,
                                    "keys": result.keys})

            def decode(self, ids, skip_special_tokens=True):
                return "Thought: test\nAction: DONE()"

            def encode(self, text, add_special_tokens=False):
                return [min(ord(c), 199) for c in text[:20]]

        return TinyVLM(), TinyProcessor()

    @staticmethod
    def _make_mock_adapter():
        """Mock WAA adapter that returns a fake screenshot."""
        from PIL import Image

        adapter = MagicMock()
        # Create a real PNG screenshot
        img = Image.new("RGB", (100, 100), color=(128, 128, 128))
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        screenshot = buf.getvalue()

        from openadapt_evals.adapters.base import (
            BenchmarkObservation, BenchmarkResult, BenchmarkTask,
        )

        adapter.observe.return_value = BenchmarkObservation(
            screenshot=screenshot, raw_observation={},
        )
        adapter.reset.return_value = BenchmarkObservation(
            screenshot=screenshot, raw_observation={},
        )
        adapter.step.return_value = (
            BenchmarkObservation(screenshot=screenshot, raw_observation={}),
            True,  # done after 1 step
            {},
        )
        adapter.load_task.return_value = BenchmarkTask(
            task_id="test-001", instruction="Test task", domain="desktop",
        )
        adapter.evaluate.return_value = BenchmarkResult(
            task_id="test-001", success=False, score=0.0,
        )
        adapter.config = MagicMock(server_url="http://localhost:5001")
        return adapter

    def test_generation_sees_pixel_values(self):
        """The model sees pixel_values during rollout generation."""
        torch = pytest.importorskip("torch")
        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        model, processor = self._make_tiny_vlm_and_processor()
        wrapper = VLMModelWrapper(model)

        # Simulate what generate_fn does
        from openadapt_evals.training.standalone.prompt import build_agent_messages
        from PIL import Image

        img = Image.new("RGB", (100, 100), color=(128, 128, 128))
        messages = build_agent_messages("Test task", include_image=True)
        text_input = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        inputs = processor(text=[text_input], images=[img], return_tensors="pt")

        # Cache vision inputs (what generate_fn does)
        wrapper.cache_vision_inputs(dict(inputs.items()))

        # Generate (what generate_fn does)
        outputs = wrapper.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs.get("pixel_values"),
            max_new_tokens=5, do_sample=True,
        )

        assert model._saw_pixel_values, (
            "Model did not see pixel_values during generation. "
            "The model is blind — this produces garbage output."
        )

    def test_trl_forward_gets_cached_pixel_values(self):
        """TRL's training forward pass gets pixel_values via the wrapper."""
        torch = pytest.importorskip("torch")
        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        model, processor = self._make_tiny_vlm_and_processor()
        wrapper = VLMModelWrapper(model)

        # Step 1: Rollout generation — cache vision inputs
        from PIL import Image
        img = Image.new("RGB", (100, 100), color=(128, 128, 128))
        inputs = processor(text=["test prompt"], images=[img], return_tensors="pt")
        wrapper.cache_vision_inputs(dict(inputs.items()))

        # Step 2: Simulate TRL's training forward pass (input_ids only!)
        model._saw_pixel_values = False
        wrapper.forward(input_ids=inputs["input_ids"])

        assert model._saw_pixel_values, (
            "Model did not see pixel_values during TRL's forward pass. "
            "The VLMModelWrapper failed to inject cached vision inputs. "
            "This means TRL's logprob recomputation is blind."
        )

    def test_output_format_not_garbage(self):
        """Generation produces parseable output, not # # # # #."""
        torch = pytest.importorskip("torch")
        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        model, processor = self._make_tiny_vlm_and_processor()
        wrapper = VLMModelWrapper(model)

        from PIL import Image
        img = Image.new("RGB", (100, 100), color=(128, 128, 128))

        from openadapt_evals.training.standalone.prompt import build_agent_messages
        messages = build_agent_messages("Test task", include_image=True)
        text_input = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )

        # Verify prompt does NOT contain <think>
        assert "<think>" not in text_input, (
            "Prompt contains <think> — model will enter thinking mode "
            "and produce garbage tokens instead of DSL actions."
        )

        # Verify prompt contains DSL format guidance
        assert "CLICK" in text_input or "Action:" in text_input, (
            "Prompt missing DSL format guidance (CLICK/TYPE/WAIT/DONE). "
            "Without this, the model doesn't know the expected output format."
        )

    def test_no_thinking_tokens_in_template(self):
        """Chat template does not inject <think> tags."""
        torch = pytest.importorskip("torch")

        _, processor = self._make_tiny_vlm_and_processor()

        tpl = getattr(processor, "chat_template", "") or ""
        assert "<think>" not in tpl, (
            "Processor chat_template contains <think>. This activates "
            "Qwen3.5 thinking mode which produces opaque reasoning tokens."
        )

    def test_vision_changes_logits(self):
        """Pixel_values actually change the model's logits (not ignored)."""
        torch = pytest.importorskip("torch")
        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        model, processor = self._make_tiny_vlm_and_processor()
        wrapper = VLMModelWrapper(model)

        input_ids = torch.tensor([[1, 2, 3, 4, 5]])

        # Forward without vision
        out_blind = wrapper.forward(input_ids=input_ids)
        logits_blind = out_blind.logits.detach().clone()

        # Cache vision + forward with injection
        pixel_values = torch.randn(1, 3, 10, 10)
        wrapper.cache_vision_inputs({"pixel_values": pixel_values})
        out_vision = wrapper.forward(input_ids=input_ids)
        logits_vision = out_vision.logits.detach()

        assert not torch.allclose(logits_blind, logits_vision, atol=1e-6), (
            "Logits identical with and without pixel_values. "
            "Either the model ignores vision inputs or the wrapper "
            "isn't injecting them. Training will produce zero gradient."
        )

    def test_wrapper_passes_peft_validation(self):
        """VLMModelWrapper passes TRL's PEFT/quantization validation.

        TRL checks isinstance(model, PeftModel) to verify adapters are
        attached to quantized models. The wrapper must pass this check.
        Without it: ValueError: "You cannot perform fine-tuning on purely
        quantized models. Please attach trainable adapters."
        """
        torch = pytest.importorskip("torch")

        try:
            from peft import PeftModel
        except ImportError:
            pytest.skip("peft not installed")

        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        # Create a mock PeftModel (has peft_config, active_adapter, etc.)
        model = MagicMock(spec=PeftModel)
        model.peft_config = {"default": MagicMock()}
        model.active_adapter = "default"
        model.parameters.return_value = iter([torch.zeros(1, requires_grad=True)])

        wrapper = VLMModelWrapper(model)

        # The critical check TRL performs
        assert isinstance(wrapper, PeftModel), (
            "isinstance(wrapper, PeftModel) must be True. "
            "TRL rejects quantized models without PEFT adapters."
        )
        assert hasattr(wrapper, "peft_config"), (
            "wrapper must expose peft_config for TRL validation."
        )

    def test_wrapper_preserves_trainable_parameters(self):
        """VLMModelWrapper exposes trainable parameters for the optimizer.

        TRL needs model.parameters() to set up the optimizer. The wrapper
        must delegate this to the wrapped model.
        """
        torch = pytest.importorskip("torch")
        from openadapt_evals.training.vlm_wrapper import VLMModelWrapper

        model, _ = self._make_tiny_vlm_and_processor()
        wrapper = VLMModelWrapper(model)

        # Verify parameters are accessible and trainable
        params = list(wrapper.parameters())
        assert len(params) > 0, "Wrapper must expose model parameters"
        assert any(p.requires_grad for p in params), (
            "At least some parameters must require grad for training"
        )
