# SPDX-FileCopyrightText: 2024 Arthit Suriyawongkul
# SPDX-FileType: SOURCE
# SPDX-License-Identifier: MIT

from inspect_ai.dataset import Sample

from inspect_evals.onet.onet import (
    extract_choice_letter,
    record_to_sample,
    split_numbered_list,
    thai_digit_to_arabic_digit,
)


class TestHelperFunctions:
    def test_extract_choice_letter(self):
        text = "2. Ans"
        assert extract_choice_letter(text) == "B"

    def test_record_to_sample(self):
        record = {
            "year": "2023",
            # This has to be "no" with quotes,
            # as the column name of the dataset put it lie this.
            '"no"': "14",
            "instruction": "ข้อใดถูก",
            "input": "1. หนึ่ง 2. สอง 3. สาม",
            "result": "2",
            "isAnswerable": True,
            "isMultipleChoice": True,
            "isSingleChoiceSolution": True,
        }
        sample = record_to_sample(record)
        assert isinstance(sample, Sample)
        assert sample.input == "ข้อใดถูก"
        assert sample.target == "B"  # result: "2" -> target: "B"
        assert sample.id.startswith("2023_14_")
        assert sample.metadata["year"] == "2023"

    def test_split_numbered_list(self):
        text = "เช่น 1. สาร A 2. สาร B 3. สาร C 4. สาร A และ B 5. สาร B และ C"
        lead_text, choices = split_numbered_list(text)
        assert lead_text == "เช่น"
        assert len(choices) == 5

        text = "1. 25 2. 17 3. 18"
        lead_text, choices = split_numbered_list(text)
        assert lead_text == ""
        assert len(choices) == 3

    def test_thai_digit_to_arabic_digit(self):
        assert thai_digit_to_arabic_digit("") == ""
        assert thai_digit_to_arabic_digit("พ.ศ. ๒๕๖๘") == "พ.ศ. 2568"
        assert thai_digit_to_arabic_digit("  A7 ") == "  A7 "
