import base64
from io import BytesIO
from unittest.mock import MagicMock

from PIL import Image

from wabot.screenshot import FULL_PAGE_SCREENSHOT, save_full_page


def png_bytes(width, height, color=(200, 30, 30)):
    buf = BytesIO()
    Image.new("RGB", (width, height), color).save(buf, format="PNG")
    return buf.getvalue()


class TestFirefoxPath:
    def test_registers_raw_command_and_decodes_base64(self, tmp_path):
        driver = MagicMock(name="driver")
        driver.capabilities = {"browserName": "firefox"}
        driver.command_executor._commands = {}
        driver.execute.return_value = {"value": base64.b64encode(png_bytes(50, 80)).decode()}
        out = tmp_path / "shot.png"
        assert save_full_page(driver, str(out)) is True
        # helper methods only exist on webdriver.Firefox, so the raw
        # geckodriver command must have been registered for Remote support
        assert driver.command_executor._commands[FULL_PAGE_SCREENSHOT] == (
            "GET",
            "/session/$sessionId/moz/screenshot/full",
        )
        driver.execute.assert_called_once_with(FULL_PAGE_SCREENSHOT)
        assert Image.open(out).size == (50, 80)


class TestChromiumPath:
    def test_single_viewport_page_is_saved_directly(self, tmp_path):
        driver = MagicMock(name="driver")
        driver.capabilities = {"browserName": "chrome"}
        # page fits in one viewport: total == viewport
        driver.execute_script.side_effect = lambda script: {
            "return document.body.parentNode.scrollWidth": 100,
            "return document.body.parentNode.scrollHeight": 60,
            "return document.documentElement.clientWidth": 100,
            "return window.innerHeight": 60,
        }.get(script)
        driver.get_screenshot_as_png.return_value = png_bytes(100, 60)
        out = tmp_path / "shot.png"
        assert save_full_page(driver, str(out)) is True
        assert Image.open(out).size == (100, 60)

    def test_tall_page_is_stitched_from_tiles(self, tmp_path):
        driver = MagicMock(name="driver")
        driver.capabilities = {"browserName": "chrome"}
        driver.execute_script.side_effect = lambda script: {
            "return document.body.parentNode.scrollWidth": 100,
            "return document.body.parentNode.scrollHeight": 150,  # 3 tiles of 60
            "return document.documentElement.clientWidth": 100,
            "return window.innerHeight": 60,
        }.get(script)
        driver.get_screenshot_as_png.return_value = png_bytes(100, 60)
        out = tmp_path / "shot.png"
        assert save_full_page(driver, str(out)) is True
        assert Image.open(out).size == (100, 150)
        # scrolled at least twice beyond the initial position
        scroll_calls = [
            c for c in driver.execute_script.call_args_list
            if c.args and str(c.args[0]).startswith("window.scrollTo")
        ]
        assert len(scroll_calls) >= 2