import base64
import os
import shutil
from pathlib import Path

import pytest
from PIL import Image

from inspect_evals.mmiu.utils import encode_img_to_base64, ensure_data, separate_options

# Define a test data directory
TEST_DATA_DIR = Path("test_data")

# Mock a smaller subset of DATA_FILES
TEST_DATA_FILES = ["Low-level-semantic", "High-level-sub-semantic"]


@pytest.mark.dataset_download
def test_ensure_data():
    """Test the ensure_data function to verify correct download and extraction of data files."""
    # Ensure the test directory is clean
    if TEST_DATA_DIR.exists():
        shutil.rmtree(TEST_DATA_DIR)

    TEST_DATA_DIR.mkdir(parents=True, exist_ok=True)

    # Call the function
    ensure_data(TEST_DATA_DIR, force_download=True, data_files=TEST_DATA_FILES)

    # Verify extraction
    for file in TEST_DATA_FILES:
        extracted_dir = TEST_DATA_DIR / file
        assert extracted_dir.exists() and extracted_dir.is_dir(), (
            f"Extraction failed for {file}"
        )

    print("Test passed: Files downloaded and extracted successfully.")


@pytest.fixture(scope="function")
def image_files():
    # Create temporary files for testing
    valid_image_path = "test_image.bmp"
    non_existent_image_path = "non_existent_image.bmp"
    corrupted_image_path = "corrupted_image.bmp"

    # Create a valid BMP image
    Image.new("RGB", (100, 100), color="red").save(valid_image_path, format="BMP")

    # Create a corrupted BMP file
    with open(corrupted_image_path, "wb") as f:
        f.write(b"not_an_image")

    yield {
        "valid": valid_image_path,
        "non_existent": non_existent_image_path,
        "corrupted": corrupted_image_path,
    }

    # Clean up temporary files after tests
    for file in [valid_image_path, corrupted_image_path]:
        if os.path.exists(file):
            os.remove(file)


@pytest.mark.parametrize(
    "test_name,image_path_key,expected_exception,expected_prefix",
    [
        ("Valid BMP image file", "valid", None, "data:image/png;base64,"),
        ("Non-existent BMP image file", "non_existent", FileNotFoundError, None),
        ("Corrupted BMP image file", "corrupted", IOError, None),
    ],
)
def test_encode_img_to_base64(
    test_name, image_path_key, expected_exception, expected_prefix, image_files
):
    # Get the actual file path from the fixture
    image_path = image_files[image_path_key]

    # When we expect an exception
    if expected_exception:
        with pytest.raises(expected_exception):
            encode_img_to_base64(image_path)
    # When we expect success
    else:
        result = encode_img_to_base64(image_path)
        assert result.startswith(expected_prefix), (
            f"Result does not start with expected prefix: {expected_prefix}"
        )
        # Extract the format from the data URL prefix
        import re

        format_match = re.search(r"data:image/(\w+);base64,", expected_prefix)
        assert format_match, f"Could not extract format from prefix: {expected_prefix}"
        expected_format = format_match.group(1).upper()

        # Re-interpret the result as a base64 string
        base64_string = result[len(expected_prefix) :]
        # Check that the base64 string is a valid image
        decoded_bytes = base64.b64decode(base64_string)

        # Round-trip check: verify we can load the decoded bytes as an image
        from io import BytesIO

        img = Image.open(BytesIO(decoded_bytes))
        assert img.format == expected_format, (
            f"Decoded image should be in {expected_format} format"
        )
        assert img.size == (100, 100), "Image dimensions should match the original"


def test_separate_options_basic():
    options_string = "Option 1\nOption 2\nOption 3\n"
    assert separate_options(options_string) == ["Option 1", "Option 2", "Option 3"]


def test_separate_options_with_prefix():
    options_string = "A: Option 1\nB: Option 2\nC: Option 3"
    assert separate_options(options_string) == ["Option 1", "Option 2", "Option 3"]


def test_separate_options_multiline():
    options_string = """
A: 1.1529 0.012747 244.44
0.41529 1.1943 -155.59
0.00087156 5.6224e-05 1.0092

B: 0.54693 0.20925 -108.35
-0.082341 1.1176 -236.48
-0.0006026 0.0001769 1.0001

C: 1.2108 -0.031741 47.374
0.20996 1.0345 -107.36
0.00054926 -6.3631e-06 1.0004

D: 1.7761 -0.053427 263.17
0.41751 1.5987 -329.46
0.00069677 3.1372e-05 1.0014
    """
    assert separate_options(options_string) == [
        "1.1529 0.012747 244.44\n0.41529 1.1943 -155.59\n0.00087156 5.6224e-05 1.0092",
        "0.54693 0.20925 -108.35\n-0.082341 1.1176 -236.48\n-0.0006026 0.0001769 1.0001",
        "1.2108 -0.031741 47.374\n0.20996 1.0345 -107.36\n0.00054926 -6.3631e-06 1.0004",
        "1.7761 -0.053427 263.17\n0.41751 1.5987 -329.46\n0.00069677 3.1372e-05 1.0014",
    ]


def test_separate_options_with_spurious_newline():
    options_string = """
A: A shirtless man is sunbathing near the pole vault area.

B: A shirtless guy does a pole vault.
C: A shirtless man lounges by the pool.
D: A shirtless man is sunbathing by the pool.
    """
    assert separate_options(options_string) == [
        "A shirtless man is sunbathing near the pole vault area.",
        "A shirtless guy does a pole vault.",
        "A shirtless man lounges by the pool.",
        "A shirtless man is sunbathing by the pool.",
    ]
