"""
读取Excel并将每行导出为dict。

用法示例:
	python merge.py --excel datasets/excel/your_file.xlsx
	python merge.py --excel datasets/excel/your_file.xlsx --out rows.json
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Iterable

import pandas as pd

columns_map = {
	'审查项':'review_item',
	'合同原文':'original_text',
	'建议':'suggest'
}


def _normalize_value(value: Any) -> Any:
	if pd.isna(value):
		return None
	return value


def read_excel_rows(excel_path: str | Path, sheet_name: str | int | None = 0) -> list[dict[str, Any]]:
	"""读取Excel并返回每行的dict列表。"""
	df = pd.read_excel(excel_path, sheet_name=sheet_name)
	rows = df.to_dict(orient="records")
	return [
		{columns_map.get(str(k).strip(), k): _normalize_value(v) for k, v in row.items()}
		for row in rows
	]


def read_excel_rows_from_path(
	excel_path: str | Path, sheet_name: str | int | None = 0
) -> list[dict[str, Any]]:
	"""读取单个Excel或目录下所有xlsx文件并返回每行的dict列表。"""
	path = Path(excel_path)
	if path.is_dir():
		all_rows: list[dict[str, Any]] = []
		for file_path in sorted(path.glob("*.xlsx")):
			rows = read_excel_rows(file_path, sheet_name=sheet_name)
			for row in rows:
				row["__source_file__"] = file_path.name
				row["ground_truth"] = '不合格'
			all_rows.extend(rows)
		return all_rows

	return read_excel_rows(path, sheet_name=sheet_name)


def _print_rows(rows: Iterable[dict[str, Any]]) -> None:
	for row in rows:
		print(row)


def main() -> None:
	parser = argparse.ArgumentParser(description="读取Excel并导出每行dict")
	parser.add_argument("--excel", default='datasets/excel', help="Excel文件路径或目录")
	parser.add_argument("--sheet", default=0, help="sheet名称或索引，默认0")
	parser.add_argument("--out", help="可选输出JSON文件路径")
	args = parser.parse_args()

	sheet: str | int | None
	if isinstance(args.sheet, str) and args.sheet.isdigit():
		sheet = int(args.sheet)
	else:
		sheet = args.sheet

	rows = read_excel_rows_from_path(args.excel, sheet_name=sheet)

	if args.out:
		out_path = Path(args.out)
		out_path.parent.mkdir(parents=True, exist_ok=True)
		out_path.write_text(json.dumps(rows, ensure_ascii=False, indent=2), encoding="utf-8")
		return

	# _print_rows(rows)
	print(f'len(rows)={len(rows)}')
	print(rows[0])


if __name__ == "__main__":
	main()