from __future__ import annotations

import json_repair
import argparse
import json
import os
from pathlib import Path
from typing import Any

import pandas as pd
from openai import OpenAI
from tenacity import Retrying, stop_after_attempt, wait_fixed

from merge_excel import read_excel_rows_from_path
from merge_json import apply_field_mapping, read_json_lists_from_dir

base_url = 'http://192.168.252.71:9002/v1'
model= 'Qwen2-72B-Instruct'
key = 'none'


# base_url = 'http://192.168.252.71:18088/api/v1'
# model = 'Qwen2-72B-Instruct'
# key = 'fastgpt-eSZPKebnYpOrsm2jsA6iithc81AseS7nknOjOavElp4s57HrTYSo8ZXW'

REVIEW_ITEM_MAP: dict[str, str] = {
	"发票审查-csc": "发票审查",
	"发票审查-merge": "发票审查",
	"价款内容审查-merge": "价款内容审查",
	"运输保险审查-merge": "运输保险审查",
	"变更取消责任审查-csc": "变更取消责任审查",
	"安装调试审查": "安装调试与指导审查",
	"标的物内容审查": "标的物审查",
	"技术指导审查": "安装调试与指导审查",
	"履行义务审查": "第三方审查",
	"三方贷款审查": "第三方审查",
	"三方货款审查": "第三方审查",
	"违约审查": "违约与延期审查",
	"违约条款审查": "违约与延期审查",
	"延期审查": "违约与延期审查",
	"责任规定审查": "违约与延期审查",
}


def normalize_review_item(review_item: Any) -> str:
	review_item_text = str(review_item).strip()
	mapped_review_item = REVIEW_ITEM_MAP.get(review_item_text)
	if mapped_review_item is not None:
		return mapped_review_item

	if review_item_text.endswith(".json"):
		review_item_text = review_item_text[:-5]

	mapped_review_item = REVIEW_ITEM_MAP.get(review_item_text)
	if mapped_review_item is not None:
		return mapped_review_item

	return review_item_text


def merge_export_rows(
	excel_path: str | Path,
	json_dir: str | Path,
	sheet_name: str | int | None = 0,
) -> list[dict[str, Any]]:
	excel_rows = read_excel_rows_from_path(excel_path, sheet_name=sheet_name)
	for row in excel_rows:
		row["__source_type__"] = "excel"

	json_rows = read_json_lists_from_dir(json_dir)
	json_rows = apply_field_mapping(json_rows)
	for row in json_rows:
		row["__source_type__"] = "json"

	return excel_rows + json_rows


def build_review_item_mapping_from_sheet(
	excel_path: str | Path,
	sheet_name: str = "提取审查",
) -> dict[str, Any]:
	path = Path(excel_path)
	files = sorted(path.glob("*.xlsx")) if path.is_dir() else [path]
	mapping: dict[str, Any] = {}

	for file_path in files:
		try:
			df = pd.read_excel(file_path, sheet_name=sheet_name)
		except Exception:
			continue

		if df.empty or len(df.columns) < 2:
			continue

		key_col = "审查要点" if "审查要点" in df.columns else df.columns[0]
		value_col = "审查提示词" if "审查提示词" in df.columns else df.columns[1]
  
		for _, row in df.iterrows():
			key = row.get(key_col)
			if pd.isna(key):
				continue
			value = row.get(value_col)
			if pd.isna(value):
				continue
			mapping[str(key).strip()] = value

	return mapping


def apply_review_item_mapping(rows: list[dict[str, Any]], mapping: dict[str, Any]) -> None:
	unmatched_review_items: set[str] = set()
	for row in rows:
		review_item = row.get("review_item")
		if review_item is None:
			unmatched_review_items.add("<None>")
			continue
		review_item_text = normalize_review_item(review_item)
		row["review_item"] = review_item_text
		mapped_value = mapping.get(review_item_text)
		if mapped_value is not None:
			row["prompt"] = mapped_value.replace('&#10;', '\n')
		else:
			unmatched_review_items.add(review_item_text)

	if unmatched_review_items:
		print(f"[mapping_unmatched_dedup] review_items={sorted(unmatched_review_items)}")


def group_rows_by_review_item_with_dedup(
	rows: list[dict[str, Any]],
) -> dict[str, list[dict[str, Any]]]:
	grouped_rows: dict[str, list[dict[str, Any]]] = {}
	seen_original_texts: dict[str, set[Any]] = {}

	for row in rows:
		review_item = row.get("review_item")
		review_item_key = "<None>" if review_item is None else str(review_item).strip()
		if review_item_key not in grouped_rows:
			grouped_rows[review_item_key] = []
			seen_original_texts[review_item_key] = set()

		original_text = row.get("original_text")
		original_text_key = original_text.strip() if isinstance(original_text, str) else original_text
		if original_text_key in seen_original_texts[review_item_key]:
			continue

		seen_original_texts[review_item_key].add(original_text_key)
		grouped_rows[review_item_key].append(row)

	return dict(sorted(grouped_rows.items(), key=lambda item: item[0]))


def print_grouped_review_item_stats(grouped_rows: dict[str, list[dict[str, Any]]]) -> None:
	total_rows = sum(len(group_rows) for group_rows in grouped_rows.values())
	print(f"grouped_review_item_count={len(grouped_rows)}")
	print(f"grouped_total_rows={total_rows}")
	for review_item, group_rows in grouped_rows.items():
		print(f"[group_stats] review_item={review_item}, count={len(group_rows)}")


def call_openai_chat_completion(
	client: OpenAI,
	model: str,
	system_prompt: str,
	user_text: str,
) -> str:
    extra_body = None
    if system_prompt:
        extra_body = {
			'variables': {
				'system': system_prompt 
			}
    	}
    resp = client.chat.completions.create(
		model=model,
		temperature=1.2,
		messages=[
			{"role": "system", "content": system_prompt},
			{"role": "user", "content": user_text},
		],
  		 extra_body=extra_body
	)
    content = resp.choices[0].message.content if resp.choices else None
    reasoning_content = resp.choices[0].message.reasoning_content if resp.choices else None
    if not isinstance(content, str) or not content.strip():
        raise ValueError("OpenAI响应缺少message.content")
    return reasoning_content,content


def filter_rows_by_llm_result(
	rows: list[dict[str, Any]],
	client: OpenAI,
	model: str,
	max_retries: int = 10,
) -> list[dict[str, Any]]:
	kept_rows: list[dict[str, Any]] = []
	
	for index, row in enumerate(rows, start=1):
		system_prompt = row.get("prompt")
		user_text = row.get("original_text")
		ground_truth = row.get("ground_truth")

		if not isinstance(system_prompt, str) or not system_prompt.strip():
			continue
		if not isinstance(user_text, str) or not user_text.strip():
			continue
		if ground_truth is None:
			continue

		ground_truth_text = str(ground_truth).strip()

		try:
			for attempt in Retrying(wait=wait_fixed(1), stop=stop_after_attempt(max_retries), reraise=True):
				with attempt:
					reasoning_content,content = call_openai_chat_completion(
						client=client,
						model=model,
						system_prompt=system_prompt,
						user_text=user_text,
					)
					result_text = json_repair.loads(content)
					if result_text is None:
						print(f"[llm_parse_failed] row={index}, attempt={attempt.retry_state.attempt_number}")
						raise ValueError("llm response json parse failed")
					final_result = result_text[0].get("result")
					if final_result != ground_truth_text:
						print(
							f"[llm_mismatch] row={index}, attempt={attempt.retry_state.attempt_number}, "
							f"original_text={user_text},ground_truth={ground_truth_text},result={result_text}"
						)
						raise ValueError("llm result mismatch")
					print(f"[llm_match] row={index}, attempt={attempt.retry_state.attempt_number},result={result_text}")
					row['rsp'] = f'<think>{reasoning_content}</think>{content}'
		except Exception as exc:
			row['rsp'] = ''
			print(f"[llm_drop] row={index}, error={exc}")
		finally:
			kept_rows.append(row)

	print(f"llm_filtered_rows={len(kept_rows)}/{len(rows)}")
	return kept_rows


def main() -> None:
	parser = argparse.ArgumentParser(description="合并Excel与JSON导出内容")
	parser.add_argument("--know", default="prompts/know.xlsx", help="提示词目录路径")
	parser.add_argument("--excel", default="datasets/excel", help="Excel文件路径或目录")
	parser.add_argument("--json-dir", default="datasets/json", help="JSON目录路径")
	parser.add_argument("--sheet", default=0, help="sheet名称或索引，默认0")
	parser.add_argument("--model", default=model, help="OpenAI模型名称")
	parser.add_argument("--max-retries", type=int, default=1, help="每条数据最大重试次数")
	parser.add_argument("--openai-base-url", default=base_url, help="OpenAI接口基础URL")
	parser.add_argument("--out", default="datasets/merged/merged_all.json", 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 = merge_export_rows(args.excel, args.json_dir, sheet_name=sheet)
	mapping = build_review_item_mapping_from_sheet(args.know, sheet_name="提取审查")
	apply_review_item_mapping(rows, mapping)

	api_key = os.getenv("OPENAI_API_KEY", key).strip()
	if not api_key:
		raise ValueError("缺少环境变量 OPENAI_API_KEY")
	client = OpenAI(api_key=api_key, base_url=args.openai_base_url)

	# rows = rows[:2]
	
	rows = filter_rows_by_llm_result(
		rows=rows,
		client=client,
		model=args.model,
		max_retries=max(1, args.max_retries),
	)

	grouped_rows = group_rows_by_review_item_with_dedup(rows)
	

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

	print(f"len(rows)={len(rows)}")
	print(f"mapped_review_item_count={len(mapping)}")
	print_grouped_review_item_stats(grouped_rows)
	print(f"saved to: {out_path}")


if __name__ == "__main__":
	main()
