
from __future__ import annotations

import json
from typing import Dict, List, Optional

from core.tool import tool, tool_func
from core.tools.segment_llm import LLMTool
from core.config import META_KEY
from loguru import logger

SUMMARY_SYSTEM_PROMPT = f'''
你是合同事实提取智能体（SegmentSummary）。

你的任务是：**基于给定的审查规则，从当前合同分段中提取“与该规则直接相关的客观事实”，并结构化输出。**

【核心原则】

你必须严格围绕“规则所需信息”进行提取。

---

【事实定义】

事实必须满足：
1. 可以在当前分段原文中直接找到对应表述；
2. 不得对原文进行抽象、概括或推断；
3. 不得补充未出现的主体、条件或数值；
4. 允许对原文做最小结构化拆分（例如金额、比例、期限）。

---

【规则驱动提取要求（关键）】

- 仅提取“该审查规则执行所需要的信息字段”
- 不得提取与该规则无关的信息（即使这些信息在文本中存在）
- 若规则未涉及某类信息，则不得输出对应字段
- 若规则涉及某字段但文本未出现，需显式标记为 "未明确"

---

【输出结构】

- 输出字段：facts
- facts 是一个对象
- 键必须来自【规则字段定义（rule_fields）】
- 不得使用预设通用维度（如“支付/违约责任”等）

---

【字段填充规则】

- 每个字段值必须是对象或对象列表
- 不得输出字符串作为字段值
- 字段内容必须为原文的最小结构化表达
- 不得改写原文含义

---

【缺失信息处理（非常重要）】

- 若规则要求的字段在当前分段未出现：
  → 必须输出该字段，并标记为：

  "未明确"

（用于后续审查判断）

---

【约束】

- 严禁编造信息
- 严禁推断未出现的内容
- 不得输出风险判断或解释
- 严格输出 JSON
'''

SUMMARY_USER_PROMPT = '''
【分段原文】
{segment_text}

【规则字段定义（仅提取这些字段）】
{rule_fields}

【任务】
请仅提取“当前分段中，与候选审查规则直接相关的客观事实”。

【特别要求】
- facts 的顶层 key 必须是规则 title
- 每个规则下仅保留与该规则直接相关的信息
- 若某规则在当前分段未出现关键信息，输出该规则并标记为 "未明确"
- 不得提取与规则无关的信息

输出 JSON。
'''

OUTPUT_EXAMPLE = '''
```json
{
  "facts": {
	"支付审查": {"方式": "银行转账", "时间": "验收后30日内"},
	"违约责任审查": {"违约金比例": "合同总金额的5%"}
  }
}
```
'''


@tool("segment_summary", "分段事实提取")
class SegmentSummaryTool(LLMTool):
	def __init__(self) -> None:
		super().__init__(SUMMARY_SYSTEM_PROMPT)

	@tool_func(
		{
			"type": "object",
			"properties": {
				"segment_id": {"type": "int"},
				"segment_text": {"type": "string"},
				"rules": {"type": "array", "items": {"type": "object"}},
				"party_role": {"type": "string"},
				"context_facts": {"type": "object"},
			},
			"required": ["segment_id", "segment_text", "rules"],
		}
	)
	def run(
		self,
		segment_id: int,
		segment_text: str,
		rules: List[Dict],
		party_role: str = "",
		context_facts: Optional[Dict] = None,
	) -> Dict:
		rules = rules or []
		try:
			return self.run_with_loop(
				self._summarize_async(segment_id, segment_text, rules, party_role, context_facts)
			)
		except Exception:
			return {}
		
	def _stringify_rule(self, rules: List[Dict]) -> str:
		lines = []
		for r in rules:
			id = r.get("id", "")
			rule_text = r.get("rule", "")
			lines.append(f"规则ID: {id}\n审查规则: {rule_text}\n")
		return "\n".join(lines)

	def _build_prompt(
		self,
		segment_text: str,
		rules: List[Dict],
		context_facts: Optional[Dict],
		party_role: str,
	) -> List[Dict[str, str]]:
		# 获取规则字段定义
		rule_fields = [
			r.get("summary") for r in rules
			if r.get("summary")
		]
		user_content = SUMMARY_USER_PROMPT.format(
			segment_text=segment_text,
			# party_role=party_role,
			# rules_json=self._stringify_rule(rules),
			rule_fields=json.dumps(rule_fields, ensure_ascii=False),
		) + OUTPUT_EXAMPLE
		return self.build_messages(user_content)

	async def _summarize_async(
		self,
		segment_id: int,
		segment_text: str,
		rules: List[Dict],
		party_role: str,
		context_facts: Optional[Dict],
	) -> Dict:
		msgs = self._build_prompt(segment_text, rules, context_facts, party_role)
		try:
			resp = await self.chat_async(msgs)
			# print("segment summary response:", resp)
			data = self.parse_first_json(resp)
			facts = data.get("facts") or {}
		except Exception as e:
			logger.info(f'Error in segment summary for segment {segment_id}: {e}')
			facts = {}
		facts[META_KEY] = {
			"segment_id": segment_id,
		}
		return facts
	
if __name__=='__main__':
    tool = SegmentSummaryTool()
    res = tool.run(
        segment_id=1,
        segment_text="甲方应于合同签订之日起30日内向乙方支付合同总金额的50%，余款在货物验收合格后30日内付清.",
		rules=[{"id": "R1", "title": "付款", "rule": "付款相关", "summary": "付款方式"}],
        context_facts={},
    )
    print(res)