"""数据加载器: JSONL/Excel 导入 + YAML 导出。"""

from __future__ import annotations

import json
from dataclasses import asdict
from pathlib import Path

import yaml

from promptgrad.models import TestCase


class Loader:
    """测试用例加载/导出器。"""

    def load_yaml_dir(self, dir_path: str) -> list[TestCase]:
        """从目录加载所有 YAML 测试用例。"""
        cases = []
        p = Path(dir_path)
        for f in sorted(p.glob("*.yaml")):
            with open(f, "r", encoding="utf-8") as fh:
                data = yaml.safe_load(fh)
            cases.append(
                TestCase(
                    id=data["id"],
                    input=data.get("input", {}),
                    golden=data.get("golden"),
                    tags=data.get("tags", []),
                )
            )
        return cases

    def load_jsonl(
        self,
        file_path: str,
        input_fields: list[str],
        golden_field: str | None = None,
    ) -> list[TestCase]:
        """从 JSONL 文件加载测试用例。"""
        cases = []
        idx = 0
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                idx += 1
                row = json.loads(line)
                input_data = {k: row[k] for k in input_fields if k in row}
                golden = row.get(golden_field) if golden_field else None
                cases.append(
                    TestCase(
                        id=f"case_{idx:03d}",
                        input=input_data,
                        golden=golden,
                        tags=[],
                    )
                )
        return cases

    def load_excel(
        self,
        file_path: str,
        input_fields: list[str],
        golden_field: str | None = None,
    ) -> list[TestCase]:
        """从 Excel 文件加载测试用例。"""
        import openpyxl

        wb = openpyxl.load_workbook(file_path, read_only=True)
        ws = wb.active
        rows = list(ws.iter_rows(values_only=True))
        if not rows:
            return []
        headers = [str(h) for h in rows[0]]
        cases = []
        for idx, row in enumerate(rows[1:], 1):
            row_dict = dict(zip(headers, row))
            input_data = {k: row_dict[k] for k in input_fields if k in row_dict}
            golden = row_dict.get(golden_field) if golden_field else None
            cases.append(
                TestCase(
                    id=f"case_{idx:03d}",
                    input=input_data,
                    golden=str(golden) if golden is not None else None,
                    tags=[],
                )
            )
        wb.close()
        return cases

    def export_yaml(self, cases: list[TestCase], output_dir: str) -> None:
        """将测试用例导出为 YAML 文件。"""
        p = Path(output_dir)
        p.mkdir(parents=True, exist_ok=True)
        for case in cases:
            file_path = p / f"{case.id}.yaml"
            data = asdict(case)
            with open(file_path, "w", encoding="utf-8") as f:
                yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
