import tempfile
import unittest
from pathlib import Path

from PIL import Image

from cockpit_agent.grounding.adapter import GroundingAdapter


class FakeGroundingModel:
    def __init__(self, output: str) -> None:
        self.output = output
        self.prompt: str | None = None

    def generate(
        self,
        image_path: str,
        prompt: str,
        max_new_tokens: int = 128,
    ) -> str:
        del image_path, max_new_tokens
        self.prompt = prompt
        return self.output


class GroundingAdapterTest(unittest.TestCase):
    def test_uses_model_bbox_without_postprocessing(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            image_path = Path(directory) / "generic.png"
            Image.new("RGB", (200, 100), color="white").save(image_path)
            model = FakeGroundingModel('{"bbox_2d": [100, 200, 400, 600]}')

            result = GroundingAdapter(model).ground(
                str(image_path),
                "车辆空调内循环控制",
            )

        self.assertEqual(result.bbox_relative, (100.0, 200.0, 400.0, 600.0))
        self.assertEqual(result.bbox_pixel, (20, 20, 80, 60))
        self.assertEqual(result.center_pixel, (50, 40))
        self.assertIn("车辆空调内循环控制", model.prompt or "")
        self.assertEqual(
            set(result.to_dict()),
            {
                "semantic_query",
                "bbox_relative",
                "bbox_pixel",
                "center_pixel",
                "image_width",
                "image_height",
                "raw_output",
            },
        )

    def test_invalid_model_output_is_not_repaired(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            image_path = Path(directory) / "generic.png"
            Image.new("RGB", (32, 32), color="white").save(image_path)
            adapter = GroundingAdapter(FakeGroundingModel("not a bbox"))
            with self.assertRaises(ValueError):
                adapter.ground(str(image_path), "任意语义控件")


if __name__ == "__main__":
    unittest.main()