import tempfile
import unittest
from pathlib import Path

from PIL import Image

from cockpit_agent.grounding.adapter import GroundingResult
from cockpit_agent.perception.roi import crop_grounding_roi, grounding_to_source


def grounding(
    *,
    bbox_pixel: tuple[int, int, int, int],
    width: int,
    height: int,
    query: str = "semantic function",
) -> GroundingResult:
    x1, y1, x2, y2 = bbox_pixel
    return GroundingResult(
        semantic_query=query,
        bbox_relative=(
            x1 / width * 1000,
            y1 / height * 1000,
            x2 / width * 1000,
            y2 / height * 1000,
        ),
        bbox_pixel=bbox_pixel,
        center_pixel=(round((x1 + x2) / 2), round((y1 + y2) / 2)),
        image_width=width,
        image_height=height,
        raw_output="model output",
    )


class ROITest(unittest.TestCase):
    def test_function_bbox_is_padded_and_cropped(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.png"
            output = root / "roi.jpg"
            Image.new("RGB", (200, 100), color="white").save(source)
            result = crop_grounding_roi(
                image_path=str(source),
                function_grounding=grounding(
                    bbox_pixel=(50, 20, 150, 80),
                    width=200,
                    height=100,
                ),
                output_path=str(output),
                padding_ratio=0.10,
            )

            with Image.open(output) as roi_image:
                roi_size = roi_image.size

        self.assertEqual(result.bbox_pixel, (40, 14, 160, 86))
        self.assertEqual((result.offset_x, result.offset_y), (40, 14))
        self.assertEqual((result.width, result.height), (120, 72))
        self.assertEqual(roi_size, (120, 72))

    def test_roi_grounding_is_converted_to_source_coordinates(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.png"
            Image.new("RGB", (200, 100), color="white").save(source)
            roi = crop_grounding_roi(
                image_path=str(source),
                function_grounding=grounding(
                    bbox_pixel=(50, 20, 150, 80),
                    width=200,
                    height=100,
                ),
                output_path=str(root / "roi.jpg"),
                padding_ratio=0.10,
            )
            roi_grounding = grounding(
                bbox_pixel=(10, 6, 30, 26),
                width=roi.width,
                height=roi.height,
                query="semantic subcontrol",
            )
            converted = grounding_to_source(roi_grounding, roi)

        self.assertEqual(converted.bbox_pixel, (50, 20, 70, 40))
        self.assertEqual(converted.center_pixel, (60, 30))
        self.assertEqual((converted.image_width, converted.image_height), (200, 100))


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