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

    def test_tall_page_stitches_tiles_bottom_aligned(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)  # window.scrollTo(...) -> None
        driver.get_screenshot_as_png.side_effect = [
            png_bytes(100, 60, (255, 0, 0)),  # tile @ y=0
            png_bytes(100, 60, (0, 255, 0)),  # tile @ y=60
            png_bytes(100, 60, (0, 0, 255)),  # tile @ y=120, clamped to paste_y=90
        ]
        out = tmp_path / "shot.png"
        assert save_full_page(driver, str(out)) is True
        img = Image.open(out)
        assert img.size == (100, 150)
        assert img.getpixel((50, 30)) == (255, 0, 0)  # first tile
        assert img.getpixel((50, 75)) == (0, 255, 0)  # second tile (rows 60-89)
        assert img.getpixel((50, 140)) == (0, 0, 255)  # third tile, bottom-aligned

    def test_wide_page_clamps_last_column(self, tmp_path):
        driver = MagicMock(name="driver")
        driver.capabilities = {"browserName": "chrome"}
        driver.execute_script.side_effect = lambda script: {
            "return document.body.parentNode.scrollWidth": 150,  # 2 columns of 100
            "return document.body.parentNode.scrollHeight": 60,
            "return document.documentElement.clientWidth": 100,
            "return window.innerHeight": 60,
        }.get(script)
        driver.get_screenshot_as_png.side_effect = [
            png_bytes(100, 60, (255, 0, 0)),  # column @ x=0
            png_bytes(100, 60, (0, 0, 255)),  # column @ x=100, clamped to paste_x=50
        ]
        out = tmp_path / "shot.png"
        assert save_full_page(driver, str(out)) is True
        img = Image.open(out)
        assert img.size == (150, 60)
        assert img.getpixel((25, 30)) == (255, 0, 0)  # left column (cols 0-49)
        assert img.getpixel((140, 30)) == (0, 0, 255)  # right column, clamped (cols 50-149)