"""数据加载器测试。"""

import json

import yaml

from promptgrad.loader import Loader


class TestLoadYAML:
    def test_load_single_case(self, tmp_path):
        case_file = tmp_path / "case_001.yaml"
        case_file.write_text(
            yaml.dump(
                {
                    "id": "case_001",
                    "input": {"query": "python list"},
                    "golden": "Python list tutorial",
                    "tags": ["basic"],
                }
            )
        )
        loader = Loader()
        cases = loader.load_yaml_dir(str(tmp_path))
        assert len(cases) == 1
        assert cases[0].id == "case_001"
        assert cases[0].golden == "Python list tutorial"

    def test_load_multiple_cases(self, tmp_path):
        for i in range(3):
            f = tmp_path / f"case_{i:03d}.yaml"
            f.write_text(yaml.dump({"id": f"case_{i:03d}", "input": {"q": f"q{i}"}, "tags": []}))
        loader = Loader()
        cases = loader.load_yaml_dir(str(tmp_path))
        assert len(cases) == 3


class TestLoadJSONL:
    def test_load_jsonl(self, tmp_path):
        jsonl_file = tmp_path / "data.jsonl"
        lines = [
            json.dumps({"query": "python list", "expected": "Python list tutorial"}),
            json.dumps({"query": "java map", "expected": "Java HashMap guide"}),
        ]
        jsonl_file.write_text("\n".join(lines))
        loader = Loader()
        cases = loader.load_jsonl(
            str(jsonl_file),
            input_fields=["query"],
            golden_field="expected",
        )
        assert len(cases) == 2
        assert cases[0].input == {"query": "python list"}
        assert cases[0].golden == "Python list tutorial"
        assert cases[0].id == "case_001"

    def test_load_jsonl_no_golden(self, tmp_path):
        jsonl_file = tmp_path / "data.jsonl"
        jsonl_file.write_text(json.dumps({"query": "test"}) + "\n")
        loader = Loader()
        cases = loader.load_jsonl(str(jsonl_file), input_fields=["query"])
        assert cases[0].golden is None


class TestLoadExcel:
    def test_load_excel(self, tmp_path):
        import openpyxl

        wb = openpyxl.Workbook()
        ws = wb.active
        ws.append(["query", "context", "expected"])
        ws.append(["python list", "beginner", "Python list tutorial"])
        ws.append(["java map", "intermediate", "Java HashMap guide"])
        excel_path = tmp_path / "data.xlsx"
        wb.save(str(excel_path))

        loader = Loader()
        cases = loader.load_excel(
            str(excel_path),
            input_fields=["query", "context"],
            golden_field="expected",
        )
        assert len(cases) == 2
        assert cases[0].input == {"query": "python list", "context": "beginner"}
        assert cases[1].golden == "Java HashMap guide"


class TestExportYAML:
    def test_export(self, tmp_path):
        from promptgrad.models import TestCase

        cases = [
            TestCase(id="case_001", input={"q": "test"}, golden="answer", tags=["t1"]),
        ]
        loader = Loader()
        loader.export_yaml(cases, str(tmp_path))
        exported = tmp_path / "case_001.yaml"
        assert exported.exists()
        data = yaml.safe_load(exported.read_text())
        assert data["id"] == "case_001"
        assert data["golden"] == "answer"
