import os

import numpy as np
import pandas as pd

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmarks.singlestep.common_models import Operation
from screensuite.benchmarks.singlestep.mmind2web.benchmark import (
    MMind2WebBenchmark,
    MMind2WebConfig,
)
from screensuite.benchmarks.singlestep.mmind2web.models import BoundingBox


def test_calculate_weighted_metrics():
    """Test the _calculate_weighted_metrics method with various scenarios."""
    # Create benchmark instance
    config = MMind2WebConfig(hf_repo="osunlp/Multimodal-Mind2Web", max_step_memory=3)
    benchmark = MMind2WebBenchmark("test_benchmark", config, ["test_domain", "test_task"])
    benchmark.datasets = {"test_domain": None, "test_task": None}  # type: ignore

    # Test case 1: Normal case with multiple splits
    accuracy_scores = {
        "test_domain_click_acc": [1.0, 0.5, 0.75],
        "test_domain_type_acc": [1.0, 0.5, 0.75],
        "test_domain_average_acc": [1.0, 0.5, 0.75],
        "test_task_click_acc": [0.8, 0.9],
        "test_task_type_acc": [0.8, 0.9],
        "test_task_average_acc": [0.8, 0.9],
    }
    metrics, ref_field = benchmark._calculate_weighted_metrics(accuracy_scores)

    # Check reference field
    assert ref_field == "weighted_average_acc"

    # Check individual split metrics
    assert np.isclose(metrics["test_domain_click_acc"], 0.75)
    assert np.isclose(metrics["test_domain_type_acc"], 0.75)
    assert np.isclose(metrics["test_domain_average_acc"], 0.75)
    assert np.isclose(metrics["test_task_click_acc"], 0.85)
    assert np.isclose(metrics["test_task_type_acc"], 0.85)
    assert np.isclose(metrics["test_task_average_acc"], 0.85)

    # Check weighted metrics
    assert np.isclose(metrics["weighted_click_acc"], 0.79)  # (0.75*3 + 0.85*2)/5
    assert np.isclose(metrics["weighted_type_acc"], 0.79)  # (0.75*3 + 0.85*2)/5
    assert np.isclose(metrics["weighted_average_acc"], 0.79)  # (0.75*3 + 0.85*2)/5

    # Test case 2: Empty scores
    empty_scores: dict[str, list[float]] = {
        "test_domain_click_acc": [],
        "test_domain_type_acc": [],
        "test_domain_average_acc": [],
        "test_task_click_acc": [],
        "test_task_type_acc": [],
        "test_task_average_acc": [],
    }

    metrics, ref_field = benchmark._calculate_weighted_metrics(empty_scores)

    # All metrics should be 0.0 for empty scores
    assert all(v == 0.0 or np.isnan(v) for v in metrics.values())

    # Test case 3: With NaN values in type accuracy
    nan_scores = {
        "test_domain_click_acc": [1.0, 0.5],
        "test_domain_type_acc": [1.0, np.nan],  # One NaN value
        "test_domain_average_acc": [1.0, 0.5],
        "test_task_click_acc": [0.8],
        "test_task_type_acc": [np.nan],  # All NaN values
        "test_task_average_acc": [0.8],
    }
    metrics, _ = benchmark._calculate_weighted_metrics(nan_scores)

    # Check metrics with NaN handling
    assert np.isclose(metrics["test_domain_click_acc"], 0.75)
    assert np.isclose(metrics["test_domain_type_acc"], 1.0)  # Only non-NaN value
    assert np.isclose(metrics["test_domain_average_acc"], 0.75)
    assert np.isclose(metrics["test_task_click_acc"], 0.8)
    assert np.isnan(metrics["test_task_type_acc"])  # All NaN values
    assert np.isclose(metrics["test_task_average_acc"], 0.8)


def test_get_annotated_input_from_sample():
    """Test the _get_annotated_input_from_sample method using the sample.parquet file."""
    # Load the sample data
    current_dir = os.path.dirname(os.path.abspath(__file__))
    sample_file = os.path.join(current_dir, "sample.parquet")
    df = pd.read_parquet(sample_file)

    # Convert the first row to a sample dictionary
    sample = {
        "operation": df["operation"].tolist(),
        "screenshot": df["screenshot"].tolist(),
        "pos_candidates": df["pos_candidates"].tolist(),
        "neg_candidates": df["neg_candidates"].tolist(),
        "target_action_index": df["target_action_index"].tolist(),
        "confirmed_task": df["confirmed_task"].tolist(),
    }

    # Create benchmark instance
    config = MMind2WebConfig(hf_repo="osunlp/Multimodal-Mind2Web", max_step_memory=3)
    benchmark = MMind2WebBenchmark("test_benchmark", config, ["test"])

    # Transform the sample
    transformed_samples = list(benchmark._get_annotated_input_from_sample(sample, EvaluationConfig(run_name="test")))

    # Verify the transformed samples
    assert len(transformed_samples) == 7

    for annotated_input in transformed_samples:
        query, (operations, bbox, _) = annotated_input.messages, annotated_input.ground_truth
        # Verify query structure
        assert isinstance(query, list), "Query should be a list of messages"
        assert len(query) == 1, "Should have one message"
        message = query[0]
        assert isinstance(message, dict), "Message should be a dictionary"
        assert message["role"] == "user", "Message role should be user"

        # Verify content structure
        content = message["content"]
        assert len([c for c in content if c["type"] == "image"]) <= benchmark.config.max_step_memory + 1, (
            f"Message should have {benchmark.config.max_step_memory + 1} or less images"
        )

        # Verify operations and bounding box
        assert isinstance(operations, list), "Operations should be a list"
        assert all(isinstance(operation, Operation) for operation in operations), (
            "All operations should be Operation instances"
        )
        assert isinstance(bbox, BoundingBox), "Bounding box should be a BoundingBox instance"

        # Verify that the number of operations matches the operation type
        if any(operation.type == "type" for operation in operations):
            assert len(operations) == 2, "Type operation should have 2 operations (click and type)"
        else:
            assert len(operations) == 1, "Non-type operation should have 1 operation (click)"
