from inspect_evals.assistant_bench.helper import (
    calculate_json_score,
    calculate_number_score,
    extract_number,
    parse_and_classify,
    parse_numeric_values,
    score_answer,
)


class TestScoreAnswer:
    def test_numeric_comparison(self):
        # Test exact matches
        assert score_answer(100, 100) == 1.0
        assert score_answer("100", "100") == 1.0

        # Test close numbers
        assert score_answer(95, 100) > 0.8
        assert score_answer("95", "100") > 0.8

        # Test different magnitudes
        assert score_answer(1000, 100) < 0.5
        assert score_answer(10, 100) < 0.5

        # Test negative numbers
        assert score_answer(-100, -100) == 1.0
        assert score_answer(-95, -100) > 0.8

    def test_string_comparison(self):
        # Test exact matches
        assert score_answer("hello world", "hello world") == 1.0

        # Test partial matches
        assert 0 < score_answer("hello", "hello world") < 1

        # Test case insensitivity
        assert score_answer("Hello World", "hello world") == 1

        # Test completely different strings
        assert score_answer("hello", "goodbye") == 0

    def test_list_comparison(self):
        # Test string lists
        assert score_answer(["hello", "world"], ["hello", "world"]) == 1.0

        # Test different order
        assert score_answer(["world", "hello"], ["hello", "world"]) == 1.0

        # Test lists of dictionaries
        assert score_answer([{"a": 1}, {"b": 2}], [{"a": 1}, {"b": 2}]) == 1.0

        # Test partial matches
        assert 0 < score_answer([{"a": 1}, {"b": 3}], [{"a": 1}, {"b": 2}]) < 1

    def test_dict_comparison(self):
        # Test simple dictionaries
        assert score_answer({"a": 1, "b": 2}, {"a": 1, "b": 2}) == 1.0

        # Test nested dictionaries
        assert (
            score_answer({"a": {"x": 1}, "b": {"y": 2}}, {"a": {"x": 1}, "b": {"y": 2}})
            == 1.0
        )

        # Test mixed types
        assert (
            score_answer(
                {"a": 1, "b": "hello", "c": [1, 2]}, {"a": 1, "b": "hello", "c": [1, 2]}
            )
            == 1.0
        )

    def test_type_mismatches(self):
        # Test types that are compatible
        assert score_answer(100, "100") == 1

        # Test different types
        assert score_answer([1, 2], {"1": 2}) == 0
        assert score_answer({"a": 1}, [1]) == 0

        # Test None values
        assert score_answer(None, 100) == 0
        assert score_answer(100, None) == 0


class TestCalculateJsonScore:
    def test_simple_json(self):
        # Test exact matches
        assert calculate_json_score({"a": 1, "b": 2}, {"a": 1, "b": 2}) == 1.0

        # Test numeric differences
        assert 0 < calculate_json_score({"a": 95, "b": 2}, {"a": 100, "b": 2}) < 1

        # Test missing keys
        assert calculate_json_score({"a": 1}, {"a": 1, "b": 2}) < 1

    def test_nested_json(self):
        # Test nested dictionaries
        assert (
            calculate_json_score(
                {"a": {"x": 1}, "b": {"y": 2}}, {"a": {"x": 1}, "b": {"y": 2}}
            )
            == 1.0
        )

        # Test nested differences
        assert (
            0
            < calculate_json_score(
                {"a": {"x": 95}, "b": {"y": 2}}, {"a": {"x": 100}, "b": {"y": 2}}
            )
            < 1
        )

    def test_mixed_types(self):
        # Test different value types
        pred = {"number": 100, "string": "hello", "list": [1, 2, 3], "dict": {"x": 1}}
        gold = {"number": 100, "string": "hello", "list": [1, 2, 3], "dict": {"x": 1}}
        assert calculate_json_score(pred, gold) == 1.0

        # Test partial matches with mixed types
        pred_partial = {
            "number": 95,
            "string": "hello world",
            "list": [1, 2],
            "dict": {"x": 2},
        }
        assert 0 < calculate_json_score(pred_partial, gold) < 1

    def test_edge_cases(self):
        # Test empty dictionaries
        assert calculate_json_score({}, {}) == 1.0

        # Test single key dictionaries
        assert calculate_json_score({"a": 1}, {"a": 1}) == 1.0

        # Test mismatched types in values
        assert (
            calculate_json_score({"a": 1, "b": "2"}, {"a": 1, "b": 2}) == 0.5
        )  # Only 'a' matches perfectly

        # Test very nested structures
        deep_dict = {"l1": {"l2": {"l3": {"value": 100}}}}
        assert calculate_json_score(deep_dict, deep_dict) == 1.0


class TestCalculateNumberScore:
    def test_exact_matches(self):
        assert calculate_number_score(100, 100) == 1.0
        assert calculate_number_score(0, 0) == 1.0
        assert calculate_number_score(-50, -50) == 1.0

    def test_close_numbers(self):
        # Numbers within same order of magnitude
        assert calculate_number_score(95, 100) > 0.9
        assert calculate_number_score(105, 100) > 0.9
        assert calculate_number_score(90, 100) > 0.8

    def test_different_magnitudes(self):
        # Test numbers of different orders of magnitude
        assert calculate_number_score(1000, 100) < 0.5
        assert calculate_number_score(10, 100) < 0.5
        assert calculate_number_score(1, 1000) < 0.5

    def test_negative_numbers(self):
        # Test handling of negative numbers
        assert calculate_number_score(-100, -100) == 1.0
        assert calculate_number_score(-95, -100) > 0.9
        assert calculate_number_score(-1000, -100) < 0.5

        # Test mixed positive and negative
        score_pos_neg = calculate_number_score(100, -100)
        score_neg_pos = calculate_number_score(-100, 100)
        assert score_pos_neg == score_neg_pos  # Should handle symmetrically

    def test_zero_handling(self):
        # Test comparisons involving zero
        assert calculate_number_score(0, 0) == 1.0
        assert calculate_number_score(0.1, 0) < 1.0
        assert calculate_number_score(0, 0.1) < 1.0

    def test_edge_cases(self):
        # Test very small numbers
        assert calculate_number_score(0.0001, 0.0002) > 0.5

        # Test very large numbers
        assert calculate_number_score(1e6, 1e6) == 1.0
        assert calculate_number_score(1e6, 1e7) < 1.0

        # Test infinity
        assert calculate_number_score(float("inf"), float("inf")) == 1.0


class TestExtractNumber:
    def test_simple_numbers(self):
        # Test integers and floats
        assert extract_number("123") == 123.0
        assert extract_number("123.45") == 123.45
        assert extract_number("-123.45") == -123.45
        assert extract_number(".45") == 0.45

    def test_currency_formats(self):
        # Test various currency formats
        assert extract_number("$123.45") == 123.45
        assert extract_number("€123.45") == 123.45
        assert extract_number("£123.45") == 123.45
        assert extract_number("123.45€") == 123.45

        # Test with thousand separators
        assert extract_number("$1,234.56") == 1234.56
        assert extract_number("1,234.56€") == 1234.56
        assert extract_number("1 234.56") == 1234.56

    def test_scientific_notation(self):
        # Test scientific notation formats
        assert extract_number("1.23e5") == 123000.0
        assert extract_number("1.23E-5") == 0.0000123
        assert extract_number("-1.23e5") == -123000.0

    def test_direct_number_inputs(self):
        # Test direct number inputs
        assert extract_number(123) == 123.0
        assert extract_number(123.45) == 123.45
        assert extract_number(-123.45) == -123.45

    def test_units_and_suffixes(self):
        # Test numbers with units
        assert extract_number("123km") == 123.0
        assert extract_number("123 km") == 123.0
        assert extract_number("123.45%") == 123.45
        assert extract_number("-123.45ms") == -123.45


class TestParseNumericValues:
    def test_simple_values(self):
        # Test simple numeric strings
        assert parse_numeric_values("123") == 123.0
        assert parse_numeric_values("123.45") == 123.45
        assert parse_numeric_values("abc") == "abc"

    def test_dict_parsing(self):
        # Test dictionary parsing
        input_dict = {
            "number": "123.45",
            "text": "abc",
            "currency": "$123.45",
            "nested": {"value": "456.78"},
        }
        expected = {
            "number": 123.45,
            "text": "abc",
            "currency": 123.45,
            "nested": {"value": 456.78},
        }
        assert parse_numeric_values(input_dict) == expected

    def test_list_parsing(self):
        # Test list parsing
        input_list = ["123.45", "abc", "$678.90", ["456.78"]]
        expected = [123.45, "abc", 678.90, [456.78]]
        assert parse_numeric_values(input_list) == expected

    def test_nested_structures(self):
        # Test complex nested structures
        input_data = {
            "values": [{"price": "$123.45"}, {"price": "678.90€"}],
            "metadata": {"count": "1000", "unit": "USD"},
        }
        expected = {
            "values": [{"price": 123.45}, {"price": 678.90}],
            "metadata": {"count": 1000.0, "unit": "USD"},
        }
        assert parse_numeric_values(input_data) == expected

    def test_non_numeric_values(self):
        # Test handling of non-numeric values
        assert parse_numeric_values(None) is None
        assert parse_numeric_values(True) is True
        assert parse_numeric_values("abc") == "abc"
        assert parse_numeric_values(["abc", "def"]) == ["abc", "def"]


class TestParseAndClassify:
    def test_json_parsing(self):
        # Test JSON string parsing
        json_str = '{"number": "123.45", "text": "abc"}'
        expected = {"number": 123.45, "text": "abc"}
        assert parse_and_classify(json_str) == expected

        # Test JSON array
        json_array = '[123.45, "abc", {"value": "678.90"}]'
        expected = [123.45, "abc", {"value": 678.90}]
        assert parse_and_classify(json_array) == expected

    def test_python_literal_parsing(self):
        # Test Python literal parsing
        # Note: String input should be exactly as it would appear in Python code
        tuple_str = "(123.45, 'abc')"
        expected = (123.45, "abc")
        assert parse_and_classify(tuple_str) == expected

        # Test set literal
        set_str = "{123.45, 'abc'}"
        # Convert to list for comparison since set ordering is not guaranteed
        result = parse_and_classify(set_str)
        assert isinstance(result, set)
        assert 123.45 in result
        assert "abc" in result

    def test_number_parsing(self):
        # Test numeric string parsing
        assert parse_and_classify("123.45") == 123.45
        assert parse_and_classify("$123.45") == 123.45
        assert parse_and_classify("123.45€") == 123.45
        assert parse_and_classify("-123.45") == -123.45

    def test_multiline_parsing(self):
        # Test multiline JSON
        multiline = """
        {"value": "123.45"}
        {"value": "678.91"}
        """
        expected = [{"value": 123.45}, {"value": 678.91}]
        assert parse_and_classify(multiline) == expected

        # Test regular multiline text
        text = """
        line 1
        line 2
        """
        assert parse_and_classify(text) == ["line 1", "line 2"]

        # Test numbers as strings multiline text
        text = """
        344.2
        43.2
        """
        assert parse_and_classify(text) == [344.2, 43.2]

    def test_fallback_behavior(self):
        # Test fallback to string for unparseable content
        assert parse_and_classify("abc") == "abc"
        assert parse_and_classify("") == ""

        # Test invalid JSON
        invalid_json = '{"key": value"}'
        assert parse_and_classify(invalid_json) == invalid_json

    def test_whitespace_handling(self):
        # Test whitespace handling
        assert parse_and_classify("  123.45  ") == 123.45
        assert parse_and_classify('\t{"value": "123.45"}\n') == {"value": 123.45}
        assert parse_and_classify("   abc   ") == "abc"
