from __future__ import annotations

import json
import re
from typing import Dict, List, Literal, Optional

from core.tool import tool, tool_func
from core.tools.segment_llm import LLMTool


ROUTER_SYSTEM_PROMPT = """
你是合同分段规则路由智能体（SegmentRuleRouter）。

你的任务是：基于“当前分段文本”，从候选审查规则中选出“应执行审查”的规则项。

【路由目标】
- 仅做规则适配判断，不输出风险结论、不输出审查建议。
- 高召回优先：只要当前分段与规则存在明确相关性，就应路由命中。
- 若候选规则明显无关，则不要命中。

【判断依据】
- 以当前分段文本为主。
- 可参考上下文记忆辅助理解术语，但不得脱离当前分段文本做臆断。

【输出约束】
- 严格输出 JSON。
- 每个命中规则需给出简短 reason，说明该分段为何与规则相关。
- 若确实没有任何相关规则，返回 {"selected_items": []}。
"""


ROUTER_USER_PROMPT = """
【当前分段文本】
{segment_text}

【上下文记忆】
{context_memories_json}

【合同立场】
{party_role}

【候选审查规则】
{candidate_rules_json}

【任务】
请从候选规则中选择当前分段应执行的审查项，并输出 selected_items。
"""


SUMMARY_ROUTER_SYSTEM_PROMPT = """
你是合同分段摘要项路由智能体（SegmentSummaryRouter）。

你的任务是：基于“当前分段文本”，从候选摘要项名称中选出当前分段应提取的摘要项。

【路由目标】
- 仅做摘要项适配判断，不输出事实摘要、不输出风险结论、不输出审查建议。
- 候选摘要项只有名称，没有规则正文、触发词或其他辅助信息。
- 高召回优先：只要当前分段明显包含某个摘要项所需的信息，就应路由命中。
- 若候选摘要项明显无关，则不要命中。

【判断依据】
- 以当前分段文本为主。
- 可参考上下文记忆辅助理解术语，但不得脱离当前分段文本做臆断。

【输出约束】
- 严格输出 JSON。
- 摘要项路由只输出命中的摘要项名称，不输出其他信息。
- 若确实没有任何相关摘要项，返回 {"selected_items": []}。
"""


SUMMARY_ROUTER_USER_PROMPT = """
【当前分段文本】
{segment_text}

【上下文记忆】
{context_memories_json}

【合同立场】
{party_role}

【候选摘要项名称】
{candidate_summaries_json}

【任务】
请从候选摘要项名称中选择当前分段应提取的摘要项，并输出 selected_items。
"""


ROUTER_OUTPUT_SCHEMA = """
```json
{
  "selected_items": [
    {
      "title": "规则标题",
      "reason": "命中原因（简短）"
    }
  ]
}
```
"""


SUMMARY_ROUTER_OUTPUT_SCHEMA = """
```json
{
  "selected_items": [
    {
      "name": "摘要项名称"
    }
  ]
}
```
"""


@tool("segment_rule_router", "分段规则路由")
class SegmentRuleRouterTool(LLMTool):
    def __init__(self) -> None:
        super().__init__(ROUTER_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_memories": {"type": "array"},
                "route_by": {"type": "string"},
            },
            "required": ["segment_id", "segment_text", "rules", "party_role"],
        }
    )
    def run(
        self,
        segment_id: int,
        segment_text: str,
        rules: List[Dict],
        party_role: str,
        context_memories: Optional[List[Dict]] = None,
        route_by: Literal["rule", "summary"] = "rule",
    ) -> Dict:
        rules = rules or []
        if route_by == "summary":
            routed_summary_names = self._route_summaries(
                segment_text=segment_text,
                rules=rules,
                party_role=party_role,
                context_memories=context_memories,
            )
            return {
                "segment_id": segment_id,
                "route_by": route_by,
                "routed_rules": [],
                "routed_rule_titles": [],
                "routed_summary_names": routed_summary_names,
            }

        routed_rules = self._route_rules(
            segment_text=segment_text,
            rules=rules,
            party_role=party_role,
            context_memories=context_memories,
        )
        return {
            "segment_id": segment_id,
            "route_by": route_by,
            "routed_rules": routed_rules,
            "routed_rule_titles": [r.get("title", "") for r in routed_rules],
            "routed_summary_names": [],
        }

    def _build_candidate_rules(self, rules: List[Dict]) -> List[Dict]:
        return [
            {r.get("title", ""): r.get("rule", "")} for r in rules if r.get("title")
        ]

    def _build_candidate_summaries(self, rules: List[Dict]) -> List[str]:
        summaries: List[str] = []
        seen: set[str] = set()
        for rule in rules:
            summary = str(rule.get("summary", "")).strip()
            if not summary or summary in seen:
                continue
            summaries.append(summary)
            seen.add(summary)
        return summaries

    def _route_rules(
        self,
        segment_text: str,
        rules: List[Dict],
        party_role: str,
        context_memories: Optional[List[Dict]],
    ) -> List[Dict]:
        if not rules:
            return []

        candidates = self._build_candidate_rules(rules)
        user_content = (
            ROUTER_USER_PROMPT.format(
                segment_text=segment_text,
                context_memories_json=json.dumps(
                    context_memories or [], ensure_ascii=False
                ),
                party_role=party_role,
                candidate_rules_json=json.dumps(candidates, ensure_ascii=False),
            )
            + ROUTER_OUTPUT_SCHEMA
        )

        llm_selected: List[Dict] = []
        try:
            resp = self.run_with_loop(
                self.chat_async(self.build_messages(user_content))
            )
            data = self.parse_first_json(resp)
            llm_selected = data.get("selected_items", []) or []
        except Exception:
            llm_selected = []

        selected_titles = {
            str(item.get("title", "")).strip()
            for item in llm_selected
            if item.get("title")
        }
        selected_reasons = {
            str(item.get("title", "")).strip(): str(item.get("reason", "")).strip()
            for item in llm_selected
            if item.get("title")
        }
        trigger_titles = self._match_trigger_titles(
            segment_text=segment_text, rules=rules
        )
        merged_titles = selected_titles | trigger_titles

        if not merged_titles:
            return []

        routed_rules: List[Dict] = []
        for rule in rules:
            title = str(rule.get("title", "")).strip()
            if not title or title not in merged_titles:
                continue

            llm_reason = selected_reasons.get(title, "")
            trigger_matched = title in trigger_titles
            if llm_reason and trigger_matched:
                reason = f"llm+trigger: {llm_reason}"
            elif llm_reason:
                reason = llm_reason
            else:
                reason = "trigger matched"

            routed_rules.append(
                {
                    "id": rule.get("id", ""),
                    "title": rule.get("title", ""),
                    "level": rule.get("level", ""),
                    "rule": rule.get("rule", ""),
                    "triggers": rule.get("triggers", ""),
                    "reason": reason,
                }
            )
        return routed_rules

    def _route_summaries(
        self,
        segment_text: str,
        rules: List[Dict],
        party_role: str,
        context_memories: Optional[List[Dict]],
    ) -> List[str]:
        candidates = self._build_candidate_summaries(rules)
        if not candidates:
            return []

        user_content = (
            SUMMARY_ROUTER_USER_PROMPT.format(
                segment_text=segment_text,
                context_memories_json=json.dumps(
                    context_memories or [], ensure_ascii=False
                ),
                party_role=party_role,
                candidate_summaries_json=json.dumps(candidates, ensure_ascii=False),
            )
            + SUMMARY_ROUTER_OUTPUT_SCHEMA
        )

        llm_selected: List[Dict] = []
        try:
            resp = self.run_with_loop(
                self.chat_async(
                    [
                        {"role": "system", "content": SUMMARY_ROUTER_SYSTEM_PROMPT},
                        {"role": "user", "content": user_content},
                    ]
                )
            )
            data = self.parse_first_json(resp)
            llm_selected = data.get("selected_items", []) or []
        except Exception:
            llm_selected = []

        selected_names = set()
        for item in llm_selected:
            name = self._selected_item_name(item)
            if name:
                selected_names.add(name)
        direct_matched_names = {
            name for name in candidates if name and name in (segment_text or "")
        }
        merged_names = selected_names | direct_matched_names
        return [name for name in candidates if name in merged_names]

    def _selected_item_name(self, item: Dict | str) -> str:
        if isinstance(item, str):
            return item.strip()
        return str(
            item.get("name")
            or item.get("summary")
            or item.get("summary_name")
            or item.get("title")
            or ""
        ).strip()

    def _match_trigger_titles(self, segment_text: str, rules: List[Dict]) -> set[str]:
        text = segment_text or ""
        matched_titles: set[str] = set()
        for r in rules:
            triggers = self._parse_triggers(str(r.get("triggers", "")))
            if triggers and any(t in text for t in triggers):
                title = str(r.get("title", "")).strip()
                if title:
                    matched_titles.add(title)
        return matched_titles

    def _parse_triggers(self, trigger_text: str) -> List[str]:
        parts = re.split(r"[，,、;；\s/|]+", trigger_text or "")
        return [p.strip() for p in parts if p.strip()]


if __name__ == "__main__":
    tool = SegmentRuleRouterTool()
    demo_rules = [
        {
            "id": "R1",
            "title": "付款触发条件明确性",
            "level": "H",
            "rule": "付款应绑定明确触发条件和验收标准。",
            "triggers": "支付,付款,验收",
        },
        {
            "id": "R2",
            "title": "违约责任对等性",
            "level": "M",
            "rule": "违约责任应当相对对等且违约金标准明确。",
            "triggers": "违约,违约金",
        },
    ]
    demo_segment_text = (
        "甲方应在合同签订后5个工作日内向乙方支付合同总价30%作为预付款，"
        "剩余70%在乙方完成交付并经甲方验收合格后30日内支付。"
        "若甲方逾期付款，每逾期一日按应付未付金额的0.05%支付违约金。"
    )

    result = tool.run(
        segment_id=1,
        segment_text=demo_segment_text,
        rules=demo_rules,
        party_role="甲方",
        context_memories=[],
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))
