from __future__ import annotations

import asyncio
import json
from typing import Dict, List, Optional

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

REVIEW_SYSTEM_PROMPT_LF = """
你是合同分段审查智能体。

基于【审查规则】审查【当前分段】。

仅审查当前分段原文，不得使用上下文或自行推断。

你必须站在“我方利益最大化”立场审查。

1. 审查原则

1.1 只输出与审查规则直接相关且证据充分的 findings。

1.2 每个 finding 必须对应一个独立判断点。

1.3 每个 finding 只能引用一个最小充分证据片段。

1.4 不得遗漏当前分段中的其他独立风险点或合规点。

1.5 result 只能为：
- "合格"
- "不合格"

1.6 当条款存在，但存在以下问题时：
- 对我方不利；
- 表述不清；
- 责任失衡；
- 条件不合理；
- 限制不足；
- 风险分配不合理；
则：
result="不合格"

1.7 当条款符合审查规则，且对我方风险可控时：
result="合格"

2. 风险等级规则
2.1 risk_level 只能为：
- "H"
- "M"
- "L"

3. suggestion 要求

3.1 当 result="不合格" 时：

suggestion 必须：
- 强化我方权利；
- 降低我方责任；
- 增加对方责任；
- 明确触发条件；
- 降低我方风险。

3.2 suggestion 不得：
- 平衡双方权利义务；
- 建议协商；
- 削弱我方优势；
- 增加我方责任；
- 减轻对方责任。

3.3 当 result="合格" 时：
suggestion="无需修改"
"""
REVIEW_USER_PROMPT_LF = """
【当前分段文本】
{segment_text}

【合同立场】
站在 {party_role} 的立场进行审查。

【审查规则】
{ruleset_text}

【任务】
请基于审查规则，仅针对当前分段文本进行审查，提取证据充分的合格条款和不合格条款，并输出 findings。

【特别要求】
- 仅基于当前分段原文进行判断，不得参考上下文、摘要或记忆信息；
- 不得扩展新的审查维度；
- 若一个段落中存在多处独立问题，必须拆分为多个 findings；
- findings 中每一项都必须包含 result；
- result 只能为：
  - "合格"
  - "不合格"

- original_text 必须：
  - 为合同原文直接引用；
  - 为最小充分证据片段；
  - 不得拼接多个不连续句子；

- 当 result="合格" 时：
  suggestion="无需修改"

- 当 result="不合格" 时：
  suggestion 必须具体、可执行，并尽量提供可直接落地的修改文本；
  若无法安全直接改写，则明确说明应补充的关键要素；

- 若 original_text 不为空，则不得认定为缺失类不合格；

- 若合同原文已经存在相关表述，但只是内容对我方不利、范围过宽、限制不足或表述不清，则不属于缺失类不合格，必须正常引用 original_text；

【输出要求】
- 若无相关或无证据充分的结论：
  {{"findings":[]}}

- 仅输出 JSON。
"""
REVIEW_OUTPUT_SCHEMA_LF = """
```json
{
  "findings": [
    {
      "rule_title": "对应审查规则标题",
      "result": "合格 或 不合格",
      "risk_level": "对应审查规则的风险等级，如 H/M/L",
      "issue": "对该条款为何合格或为何不合格的详细说明，需基于规则与案例文本解释",
      "original_text": "合同原文片段的直接引用",
      "suggestion": "若 result=合格 则填写“无需修改”；若 result=不合格 则填写可直接替换原文或新增的条款措辞"
    }
  ]
}
"""

REVIEW_SYSTEM_PROMPT_JP = """
你是一个专业的合同分段审查智能体（SegmentReview）。

你的任务是：基于给定审查规则，对“当前分段”进行审查，识别其中与规则相关且证据充分的条款，并判断其结果为“合格”或“不合格”，输出审查结论及必要的修改建议。

【审查范围】
你只能审查当前分段自身已经明确体现的内容。
你只能识别以下两类结果：
1. 合格条款：当前分段中存在与审查规则相关的明确表述，且该表述符合规则要求；
2. 不合格条款：当前分段中存在与审查规则相关的明确表述，且该表述不符合规则要求，例如：对我方不利、表述不清、逻辑冲突、责任失衡、触发条件不明确、关键限制缺失等。

【审查原则】
- 严格基于给定的审查规则进行审查，不得脱离规则自行扩展审查标准。
- 只审查当前分段原文，不得使用上下文信息补充、修正或推断当前分段含义。
- 优先识别“确定成立”的合格或不合格结论，不输出模糊怀疑类表述。

【完整性要求（非常重要）】
你必须对当前分段进行“穷举式审查”，不得只输出部分结果。

执行方式：
- 应逐句扫描当前分段
- 对每一句或关键子句，判断其是否与审查规则相关
- 只要存在证据充分的问题或合格表述，必须全部列出，不得遗漏

特别要求：
- 不得因为已找到1条或少量finding而提前停止
- 若一个段落中存在多处问题，必须分别输出多个 findings
- findings 数量应与段落中实际存在的问题数量大致一致，不得明显偏少

错误示例（禁止）：
- 一个段落有多个风险点，但只输出1条

正确行为：
- 覆盖所有可以独立成立的审查点

【结果判定规则】
- result 只能取以下两个值之一：
  - "合格"：当前分段存在与规则相关的明确内容，且符合该规则要求；
  - "不合格"：当前分段存在与规则相关的明确内容，且不符合该规则要求。
- 如果当前分段与某条审查规则无关，或虽疑似相关但证据不足，则不得生成 finding。

【证据要求】
每个 findings 都必须包含 original_text，且必须是合同原文的直接引用。

【单一证据约束（非常重要）】
每一个 finding 必须只对应一个“独立判断点”和一个“最小证据句”。

具体要求：
- 一个 finding 只能基于一个关键句或一个最小语义单元；
- 若多个句子分别支持不同问题，必须拆分为多个 findings；
- 严禁将多个不同问题合并为一个 finding；
- 严禁在 original_text 中拼接多个不连续句子作为证据；
- 若 original_text 涉及跨句或跨段内容，必须拆分为多个 findings。

判断标准：
- 如果去掉 original_text 中的一部分，仍能形成一个独立判断 → 说明应该拆分

【issue 要求】
- issue 必须说明：该条款为什么合格或为什么不合格。
- 当 result="合格" 时，issue 应说明该表述满足了什么规则要求、为什么可认定为合格。
- 当 result="不合格" 时，issue 应说明该表述违反了什么规则要求、为什么构成风险或缺陷。
- issue 必须紧扣规则和原文，不得空泛评价。

【建议要求】
- suggestion 必须具体、可执行。
- 当 result="不合格" 时：
  - 若能在当前分段内直接修正，请给出可直接替换或新增的条款措辞；
  - 若无法直接改写，请给出明确修改方向和应补充的关键要素；
  - 不得只写“建议协商”“建议完善”等空泛表述。
- 当 result="合格" 时：
  - suggestion 应简洁填写，可写“无需修改”；
  - 不得为了凑内容而提出与审查结论无关的修改建议。

【输出约束】
- 严格按照指定 JSON Schema 输出。
- 不得输出任何 JSON 之外的解释性文字。
- 若未发现证据充分的合格或不合格条款，返回 {"findings": []}。

在生成最终 JSON 之前，你必须执行以下内部步骤（不输出）：

Step A：将当前分段拆分为若干句子或语义单元
Step B：逐句判断该句是否涉及任一审查规则
Step C：若涉及规则，判断其为合格或不合格
Step D：为每一个成立的判断生成一个 finding

只有完成上述穷举后，才允许输出最终结果

提示：在合同审查中，一个分段通常可能包含多个独立风险点或合规点，findings 数量通常大于1，除非该段确实只涉及单一事项。
"""
REVIEW_USER_PROMPT_JP = """
【当前分段文本】
{segment_text}

【合同立场】
站在 {party_role} 的立场进行审查。

【审查规则】
{ruleset_text}

【任务】
请基于审查规则，参考审查案例，仅针对当前分段文本进行审查，提取证据充分的合格条款和不合格条款，并输出审查结果。

【特别要求】
- 仅基于当前分段原文进行判断，不得参考任何上下文、摘要或记忆信息。
- 若一个段落中存在多处问题，必须分别输出多个 findings。
- findings 中每一项都必须包含 result 字段，且 result 只能为 "合格" 或 "不合格"。
- findings 中的 original_text 必须为合同原文直接引用，且应为最小充分证据片段。
- 当 result="合格" 时，suggestion 填写“无需修改”。
- 当 result="不合格" 时，suggestion 应尽量提供可直接落地的修改文本；若无法安全地直接改写，请给出明确的修改方向和应补充的关键要素。
- 若无相关或无证据充分的结论，返回 {{"findings": []}}。

【输出要求】
- 仅输出 JSON。
"""
REVIEW_OUTPUT_SCHEMA_JP = """
```json
{
  "findings": [
    {
      "rule_title": "对应审查规则标题",
      "result": "合格 或 不合格",
      "issue": "对该条款为何合格或为何不合格的详细说明，需基于规则与案例文本解释",
      "original_text": "合同原文片段的直接引用",
      "suggestion": "若 result=合格 则填写“无需修改”；若 result=不合格 则填写可直接替换原文或新增的条款措辞"
    }
  ]
}
"""

LEVEL_WEIGHT = {"H": 3, "M": 2, "L": 1}


def _norm(text: str) -> str:
    if not text:
        return ""
    text = text.lower()
    text = re.sub(r"\s+", "", text)
    text = re.sub(r"[，。；：、（）()“”\"']", "", text)
    return text


def _has_evidence(f: Dict) -> bool:
    return bool(f.get("original_text"))


@tool("segment_review", "合同分段审查")
class SegmentReviewTool(LLMTool):
    def __init__(self):
        if use_lufa:
            super().__init__(REVIEW_SYSTEM_PROMPT_LF)
        else:
            super().__init__(REVIEW_SYSTEM_PROMPT_JP)

    @tool_func(
        {
            "type": "object",
            "properties": {
                "segment_id": {"type": "int"},
                "segment_text": {"type": "string"},
                "rules": {"type": "array", "items": {"type": "object"}},
                "merge_rules_prompt": {"type": "boolean"},
                "party_role": {"type": "string"},
                "context_summaries": {"type": "array"},
                "context_memories": {"type": "array"},
            },
            "required": ["segment_id", "segment_text", "rules", "party_role"],
        }
    )
    def run(
        self,
        segment_id: str,
        segment_text: str,
        rules: List[Dict],
        party_role: str,
        context_summaries: Optional[List[Dict]] = None,
        context_memories: Optional[List[Dict]] = None,
        merge_rules_prompt: bool = True,
    ) -> Dict:
        rules = rules or []
        result = self._evaluate_rules(
            party_role,
            segment_id,
            segment_text,
            rules,
            context_summaries,
            context_memories,
            merge_rules_prompt=merge_rules_prompt,
        )

        overall = "revise" if (result["findings"]) else "pass"

        return {
            "segment_id": segment_id,
            "overall_conclusion": overall,
            "findings": result["findings"],
        }

    def _stringify_rule(self, rule: Dict) -> str:
        res = ""
        res += f"## 审查项标题\n{rule.get('title','')}\n"
        res += f"## 审查规则\n{rule.get('rule','')}\n"
        res += f"## 风险等级\n{rule.get('level','')}\n"
        res += f"## 建议模板\n{rule.get('suggestion_template','')}\n"
        res += f"## 参考案例\n{rule.get('case','')}\n"
        return res

    def _build_prompt(
        self,
        party_role: str,
        rule: Dict,
        segment_id: int,
        segment_text: str,
        context_summaries: Optional[List[Dict]],
        context_memories: Optional[List[Dict]],
    ) -> List[Dict[str, str]]:
        user_content = (
            REVIEW_USER_PROMPT_JP if not use_lufa else REVIEW_USER_PROMPT_LF
        ).format(
            segment_id=segment_id,
            segment_text=segment_text,
            party_role=party_role,
            # context_facts_json=json.dumps(context_summaries or [], ensure_ascii=False),
            # context_memories_json=json.dumps(context_memories or [], ensure_ascii=False),
            ruleset_text=self._stringify_rule(rule),
        ) + (
            REVIEW_OUTPUT_SCHEMA_JP if not use_lufa else REVIEW_OUTPUT_SCHEMA_LF
        )
        return self.build_messages(user_content)

    def _build_prompt_with_rules(
        self,
        party_role: str,
        rules: List[Dict],
        segment_id: int,
        segment_text: str,
        context_summaries: Optional[List[Dict]],
        context_memories: Optional[List[Dict]],
    ) -> List[Dict[str, str]]:
        ruleset_text = "\n\n".join([self._stringify_rule(rule) for rule in rules])
        user_content = REVIEW_USER_PROMPT.format(
            segment_id=segment_id,
            segment_text=segment_text,
            party_role=party_role,
            context_facts_json=json.dumps(context_summaries or [], ensure_ascii=False),
            context_memories_json=json.dumps(
                context_memories or [], ensure_ascii=False
            ),
            ruleset_text=ruleset_text,
        ) + (REVIEW_OUTPUT_SCHEMA_JP if not use_lufa else REVIEW_OUTPUT_SCHEMA_LF)
        return self.build_messages(user_content)

    async def _evaluate_rules_async(
        self,
        party_role: str,
        segment_id: int,
        segment_text: str,
        rules: List[Dict],
        context_summaries: Optional[List[Dict]],
        context_memories: Optional[List[Dict]],
        merge_rules_prompt: bool = False,
    ) -> Dict[str, List[Dict]]:
        if merge_rules_prompt:
            msgs = [
                self._build_prompt_with_rules(
                    party_role,
                    rules,
                    segment_id,
                    segment_text,
                    context_summaries,
                    context_memories,
                )
            ]
        else:
            msgs = [
                self._build_prompt(
                    party_role,
                    rule,
                    segment_id,
                    segment_text,
                    context_summaries,
                    context_memories,
                )
                for rule in rules
            ]
        if not msgs:
            return {"findings": []}

        try:
            responses = await self.chat_batch_async(msgs)
        except Exception as e:
            logger.error(
                f"segment_review Error in segment review for segment {segment_id}: {e}"
            )
            return {"findings": []}

        all_findings: List[Dict] = []
        if merge_rules_prompt:
            data = self.parse_first_json(responses[0])
            findings = data.get("findings", []) or []
            for f in findings:
                if "rule_title" not in f:
                    f["rule_title"] = ""
                if "risk_level" not in f:
                    f["risk_level"] = ""
            all_findings.extend(findings)
        else:
            for idx, resp in enumerate(responses):
                data = self.parse_first_json(resp)
                findings = data.get("findings", []) or []
                # 为每条 finding 添加对应的 rule_title 和 risk_level
                for f in findings:
                    if "rule_title" not in f or not f["rule_title"]:
                        f["rule_title"] = rules[idx].get("title", "")
                    if "risk_level" not in f or not f["risk_level"]:
                        f["risk_level"] = rules[idx].get("level", "")
                all_findings.extend(findings)

        return {
            "findings": all_findings,
        }

    def _evaluate_rules(
        self,
        party_role: str,
        segment_id: int,
        segment_text: str,
        rules: List[Dict],
        context_summaries: Optional[List[Dict]],
        context_memories: Optional[List[Dict]],
        merge_rules_prompt: bool = False,
    ) -> Dict[str, List[Dict]]:
        try:
            return asyncio.run(
                self._evaluate_rules_async(
                    party_role,
                    segment_id,
                    segment_text,
                    rules,
                    context_summaries,
                    context_memories,
                    merge_rules_prompt=merge_rules_prompt,
                )
            )
        except RuntimeError:
            loop = asyncio.get_event_loop()
            return loop.run_until_complete(
                self._evaluate_rules_async(
                    party_role,
                    segment_id,
                    segment_text,
                    rules,
                    context_summaries,
                    context_memories,
                    merge_rules_prompt=merge_rules_prompt,
                )
            )


if __name__ == "__main__":
    tool = SegmentReviewTool()
    segment_text = """
      answer: 1.1“甲方（买方）”是指【冕宁县穗发新能源有限公司 】，包括其指定继承人（其指定继承人将全面继承需方在本合同的权利、义务和责任）。
      answer: 14.11由于买方与卖方的合同分包商和外购设备供货商没有直接的合同关系，故本合同设备的卖方的分包和外购设备的付款由卖方负责。但如果发生由于个别原因(包括但不限于买方虽按时向卖方付款而卖方没有按时向其分包商或外购设备供货商付款等情形)导致卖方的分包和外购设备有可能无法按时交货以至于影响施工进度的情况，买方有权暂时中止向卖方付款。在卖方向其分包商或外购设备供货商支付相关款项后，买方将继续向卖方付款，同时买方还将追究卖方延误工期的责任。如果卖方仍未向其分包商或外购设备供货商付款，买方将出于保障工程进度的目的，有权直接向其分包商或外购设备供货商付款。但在此情况下，卖方必须协助买方同卖方的分包商或外购设备供货商另行签订转付款协议书，同时该协议书中此转付款连同买方发生的贷款利息将从下一笔买方向卖方的应付款中扣除。
      answer: 14.12若买方认为卖方因财务或其他问题未能履行本合同内的义务，买方有权自行或另请其他方履行本合同余下的义务。卖方保证分包合同中将载有规定，在卖方无法继续经营或履行分包合同的情况下，卖方在各分包合同下的权利自动转让给买方或买方指定的其他方。
      answer: 20.8.2  卖方资质出现失效、未通过年审（年检）或被主管部门注销的，买方有权单方面解除合同并将合同未完成事项转由有资质的单位承接，已完成的事项按实结算。如因上述情形导致买方损失的，卖方应予完全赔偿。    """
    result = tool.run(
        segment_id=1,
        segment_text=segment_text,
        rules=[
            {
                "title": "第三方审查",
                "rule": """
1）货款支付不能涉及第三方（业主或委托付款方）
2）不能明确提及甲方将履行义务转移给第三方（业主或委托付款方）
3）买方转移债务到第三方（业主或委托付款方），由卖方直接向第三方（业主或委托付款方）行使债权，审查不合格
""",
            }
        ],
        party_role="麓谷发展",
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))

    # filter_result = tool.filter(1,result)
    # print(json.dumps(filter_result, ensure_ascii=False, indent=2))
