import json
from io import BytesIO

from PIL import Image

from screensuite.benchmarks.singlestep.mmind2web.models import (
    Action,
    AggregatedTraceSample,
    Mind2WebActionStep,
    Operation,
    Trace,
)


def create_test_image():
    """Create a simple test image."""
    img = Image.new("RGB", (100, 100), color="white")
    img_byte_arr = BytesIO()
    img.save(img_byte_arr, format="PNG")
    return img_byte_arr.getvalue()


def test_to_trace_normal_operation():
    """Test to_trace with normal valid data."""
    # Create test data
    operations = [json.dumps({"op": "CLICK", "original_op": "CLICK", "value": ""})]
    screenshots = [create_test_image()]
    pos_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "10,10,20,20"}'})]]
    neg_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "30,30,20,20"}'})]]
    target_action_indexes = [0]
    confirmed_tasks = ["Test task"]

    sample = AggregatedTraceSample(
        operations=operations,
        initial_screenshots=screenshots,
        pos_candidates=pos_candidates,
        neg_candidates=neg_candidates,
        target_action_indexes=target_action_indexes,
        step_tasks=confirmed_tasks,
    )

    trace = sample.to_trace()

    assert isinstance(trace, Trace)
    assert len(trace.steps) == 1
    step = trace.steps[0]
    assert step.action.operations[0].type == "click"
    assert len(step.action.operations) == 1
    assert step.bounding_box is not None


def test_to_trace_type_operation():
    """Test to_trace with TYPE operation."""
    operations = [json.dumps({"op": "TYPE", "original_op": "TYPE", "value": "test input"})]
    screenshots = [create_test_image()]
    pos_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "10,10,20,20"}'})]]
    neg_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "30,30,20,20"}'})]]
    target_action_indexes = [0]
    confirmed_tasks = ["Test task"]

    sample = AggregatedTraceSample(
        operations=operations,
        initial_screenshots=screenshots,
        pos_candidates=pos_candidates,
        neg_candidates=neg_candidates,
        target_action_indexes=target_action_indexes,
        step_tasks=confirmed_tasks,
    )

    trace = sample.to_trace()
    step = trace.steps[0]
    assert len(step.action.operations) == 2
    assert step.action.operations[0].type == "click"
    assert step.action.operations[1].type == "type"
    assert step.action.operations[1].args["text"] == "test input"


def test_to_trace_missing_screenshot():
    """Test to_trace with missing screenshot."""
    operations = [json.dumps({"op": "CLICK", "original_op": "CLICK", "value": ""})]
    screenshots = [b"\x00"]  # Missing screenshot
    pos_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "10,10,20,20"}'})]]
    neg_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "30,30,20,20"}'})]]
    target_action_indexes = [0]
    confirmed_tasks = ["Test task"]

    sample = AggregatedTraceSample(
        operations=operations,
        initial_screenshots=screenshots,
        pos_candidates=pos_candidates,
        neg_candidates=neg_candidates,
        target_action_indexes=target_action_indexes,
        step_tasks=confirmed_tasks,
    )

    trace = sample.to_trace()
    assert len(trace.steps) == 0


def test_to_trace_invalid_bounding_box():
    """Test to_trace with invalid bounding box."""
    operations = [json.dumps({"op": "CLICK", "original_op": "CLICK", "value": ""})]
    screenshots = [create_test_image()]
    pos_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "invalid"}'})]]
    neg_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "30,30,20,20"}'})]]
    target_action_indexes = [0]
    confirmed_tasks = ["Test task"]

    sample = AggregatedTraceSample(
        operations=operations,
        initial_screenshots=screenshots,
        pos_candidates=pos_candidates,
        neg_candidates=neg_candidates,
        target_action_indexes=target_action_indexes,
        step_tasks=confirmed_tasks,
    )

    trace = sample.to_trace()
    assert len(trace.steps) == 0


def test_to_trace_no_positive_candidates():
    """Test to_trace with no positive candidates."""
    actions = [Action(operations=[Operation(type="click", args={})])]
    operations = [json.dumps({"op": "CLICK", "original_op": "CLICK", "value": ""})]
    screenshots = [create_test_image()]
    pos_candidates = [[]]  # type: ignore
    neg_candidates = [[json.dumps({"attributes": '{"bounding_box_rect": "30,30,20,20"}'})]]
    target_action_indexes = [0]
    confirmed_tasks = ["Test task"]

    sample = AggregatedTraceSample(
        operations=operations,
        initial_screenshots=screenshots,
        pos_candidates=pos_candidates,
        neg_candidates=neg_candidates,
        target_action_indexes=target_action_indexes,
        step_tasks=confirmed_tasks,
    )

    trace = sample.to_trace()
    assert len(trace.steps) == 0


def test_operation_to_string():
    """Test the to_string method of Operation class."""
    # Test click action
    click_action = Operation(type="click", args={"x": 100, "y": 100})
    assert click_action.to_string() == "click(100,100)"

    # Test type action
    type_action = Operation(type="type", args={"text": "hello world"})
    assert type_action.to_string() == 'type("hello world")'


def test_step_action_to_string():
    """Test the action_to_string method of ActionStep class."""
    # Create a step with multiple actions
    step = Mind2WebActionStep(
        initial_screenshot=Image.new("RGB", (100, 100)),
        action=Action(operations=[Operation(type="click", args={"x": 100, "y": 100})]),
        pos_candidates=[],
        neg_candidates=[],
        target_action_index=0,
        task="Test task",
    )

    # Test single action
    assert step.action_to_string() == "\nAction: click(100,100)\n"

    # Test multiple actions
    step.add_operation(Operation(type="type", args={"text": "hello"}))
    assert step.action_to_string() == '\nAction: click(100,100); type("hello")\n'

    # Test empty actions
    empty_step = Mind2WebActionStep(
        initial_screenshot=Image.new("RGB", (100, 100)),
        action=Action(operations=[Operation(type="click", args={"x": 100, "y": 100})]),
        pos_candidates=[],
        neg_candidates=[],
        target_action_index=0,
        task="Test task",
    )
    assert empty_step.action_to_string() == "\nAction: click(100,100)\n"
