from __future__ import annotations

import json
import re
import unicodedata
from typing import Any, Callable, Dict, List, Optional

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

MERGER_SYSTEM_PROMPT = """
你将收到同一组 findings 的 issue 与 suggestion 列表，请做信息融合而非机械拼接。

要求：
1. 输入中已经包含同组条款原文`original_text`，请仅将其作为分析依据。
2. `issue`：提炼并合并组内风险点，去重、保留关键信息，语言精炼。
3. `suggestion`：合并为一条可执行建议，必须基于输入原文的具体表述来给出，避免空泛、泛化或与原文脱节，必要时按“先补充条款、再明确标准”这类逻辑组织。
4. 禁止输出与输入无关的信息。
"""

MERGER_USER_PROMPT = """
输入：
{payload}

输出 JSON，格式如下：
"""

OUTPUT_EXAMPLE = """
```json
{
  "issue": "提炼合并后的风险点",
  "suggestion": "提炼合并后的建议"
}
```
"""


def _as_dict(item: Any) -> Dict[str, Any]:
    if isinstance(item, dict):
        return dict(item)
    if hasattr(item, "__dict__"):
        return dict(getattr(item, "__dict__"))
    return {}


def _normalize_finding(raw: Dict[str, Any]) -> Dict[str, Any]:
    risk_level = str(raw.get("risk_level", "") or "").upper()
    if risk_level not in {"H", "M", "L", ""}:
        risk_level = ""
    result = str(raw.get("result", "") or "")
    if result not in {"合格", "不合格", ""}:
        result = ""
    return {
        "rule_title": str(raw.get("rule_title", "") or ""),
        "segment_id": int(raw.get("segment_id", 0) or 0),
        "original_text": str(raw.get("original_text", "") or ""),
        "issue": str(raw.get("issue", "") or ""),
        "risk_level": risk_level,
        "suggestion": str(raw.get("suggestion", "") or ""),
        "result": result,
    }


def _merge_unique_text(parts: List[str]) -> str:
    merged: List[str] = []
    seen = set()
    for part in parts:
        text = str(part or "").strip()
        if not text or text in seen:
            continue
        seen.add(text)
        merged.append(text)
    return "\n".join(merged)


def _max_suffix_prefix_overlap(left: str, right: str) -> int:
    max_k = min(len(left), len(right))
    for k in range(max_k, 0, -1):
        if left[-k:] == right[:k]:
            return k
    return 0


def _merge_text_union(base: str, other: str) -> str:
    left = str(base or "")
    right = str(other or "")
    if not left:
        return right
    if not right:
        return left
    if left in right:
        return right
    if right in left:
        return left

    ov1 = _max_suffix_prefix_overlap(left, right)
    ov2 = _max_suffix_prefix_overlap(right, left)
    if ov1 >= ov2 and ov1 >= 3:
        return left + right[ov1:]
    if ov2 > ov1 and ov2 >= 3:
        return right + left[ov2:]

    return f"{left}\n{right}"


def _has_substring_overlap(a: str, b: str, min_common_len: int = 5) -> bool:
    left = _normalize_text_for_match(str(a or ""))
    right = _normalize_text_for_match(str(b or ""))
    if not left or not right:
        return False
    # Only treat edge overlap as related: suffix(left)->prefix(right) or reverse.
    overlap_len = max(
        _max_suffix_prefix_overlap(left, right),
        _max_suffix_prefix_overlap(right, left),
    )
    return overlap_len >= min_common_len


def _normalize_text_for_match(text: str) -> str:
    # Normalize full-width/half-width and case, then remove whitespace and punctuation.
    normalized = unicodedata.normalize("NFKC", str(text or "")).lower()
    normalized = re.sub(r"[\s\W_]+", "", normalized, flags=re.UNICODE)
    return normalized


def _risk_rank(level: str) -> int:
    return {"": 0, "L": 1, "M": 2, "H": 3}.get(level, 0)


def _merge_group(
    group: List[Dict[str, Any]],
    field_merger: Optional[Callable[[List[Dict[str, Any]]], Dict[str, str]]] = None,
) -> Dict[str, Any]:
    if not group:
        return _normalize_finding({})

    merged_text = ""
    for item in group:
        merged_text = _merge_text_union(merged_text, item.get("original_text", ""))

    rule_titles = [item.get("rule_title", "") for item in group]
    issues = [item.get("issue", "") for item in group]
    suggestions = [item.get("suggestion", "") for item in group]
    segment_ids = [int(item.get("segment_id", 0) or 0) for item in group]
    risk_levels = [str(item.get("risk_level", "") or "") for item in group]
    results = [str(item.get("result", "") or "") for item in group]

    best_risk = ""
    for level in risk_levels:
        if _risk_rank(level) > _risk_rank(best_risk):
            best_risk = level

    if "不合格" in results:
        merged_result = "不合格"
    elif "合格" in results:
        merged_result = "合格"
    else:
        merged_result = ""

    merged_issue = _merge_unique_text(issues)
    merged_suggestion = _merge_unique_text(suggestions)
    if field_merger:
        try:
            merged_fields = field_merger(group)
            merged_issue = str(merged_fields.get("issue", "") or merged_issue)
            merged_suggestion = str(
                merged_fields.get("suggestion", "") or merged_suggestion
            )
        except Exception:
            error_msg = traceback.format_exc()
            logger.error(
                f"Error in field_merger, fallback to unique merge:{error_msg}",
                exc_info=True,
            )
            # Fall back to deterministic merge when model merge is unavailable.
            pass

    return _normalize_finding(
        {
            "rule_title": " / ".join([t for t in dict.fromkeys(rule_titles) if t]),
            "segment_id": min(segment_ids) if segment_ids else 0,
            "original_text": merged_text,
            "issue": merged_issue,
            "risk_level": best_risk,
            "suggestion": merged_suggestion,
            "result": merged_result,
        }
    )


def _rule_based_merge(
    findings: List[Dict[str, Any]],
    field_merger: Optional[Callable[[List[Dict[str, Any]]], Dict[str, str]]] = None,
) -> List[Dict[str, Any]]:
    n = len(findings)
    if n <= 1:
        return findings

    visited = [False] * n
    groups: List[List[Dict[str, Any]]] = []

    for i in range(n):
        if visited[i]:
            continue
        stack = [i]
        visited[i] = True
        group_idx: List[int] = []

        while stack:
            cur = stack.pop()
            group_idx.append(cur)
            cur_text = findings[cur].get("original_text", "")
            for j in range(n):
                if visited[j]:
                    continue
                if _has_substring_overlap(
                    cur_text, findings[j].get("original_text", "")
                ):
                    visited[j] = True
                    stack.append(j)

        groups.append([findings[idx] for idx in group_idx])

    merged_findings: List[Dict[str, Any]] = []
    for group in groups:
        if _should_skip_group_merge(group):
            merged_findings.extend([_normalize_finding(item) for item in group])
            continue
        merged_findings.append(_merge_group(group, field_merger=field_merger))
    return merged_findings


def _deterministic_field_merge(group: List[Dict[str, Any]]) -> Dict[str, str]:
    return {
        "issue": _merge_unique_text([item.get("issue", "") for item in group]),
        "suggestion": _merge_unique_text(
            [item.get("suggestion", "") for item in group]
        ),
    }


def _should_skip_group_merge(group: List[Dict[str, Any]]) -> bool:
    if len(group) <= 1:
        return True
    for item in group:
        if not str(item.get("original_text", "") or "").strip():
            return True
    return False


@tool("segment_merger", "同证据 findings 合并")
class SegmentMergerTool(LLMTool):
    def __init__(self) -> None:
        super().__init__(MERGER_SYSTEM_PROMPT)

    @tool_func(
        {
            "type": "object",
            "properties": {
                "findings": {"type": "array", "items": {"type": "object"}},
                "merge_mode": {
                    "type": "string",
                    "enum": ["llm", "rule"],
                    "default": "llm",
                },
            },
            "required": ["findings"],
        }
    )
    def run(
        self, findings: List[Dict[str, Any]], merge_mode: str = "llm"
    ) -> Dict[str, List[Dict[str, Any]]]:
        normalized_findings = self._normalize_findings(findings)
        if not normalized_findings:
            return {"findings": []}

        mode = str(merge_mode or "llm").lower()
        field_merger = self._resolve_field_merger(mode)
        merged = _rule_based_merge(normalized_findings, field_merger=field_merger)
        return {"findings": merged}

    def _normalize_findings(
        self, findings: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        return [_normalize_finding(_as_dict(item)) for item in (findings or [])]

    def _resolve_field_merger(
        self, mode: str
    ) -> Callable[[List[Dict[str, Any]]], Dict[str, str]]:
        if mode == "llm":
            return self._safe_llm_field_merge
        return _deterministic_field_merge

    def _safe_llm_field_merge(self, group: List[Dict[str, Any]]) -> Dict[str, str]:
        try:
            return self._merge_issue_suggestion_with_llm(group)
        except Exception:
            error_msg = traceback.format_exc()
            logger.error(
                f"LLM field merge failed, fallback to deterministic merge:{error_msg}",
                exc_info=True,
            )
            return _deterministic_field_merge(group)

    def _merge_issue_suggestion_with_llm(
        self, group: List[Dict[str, Any]]
    ) -> Dict[str, str]:
        payload = {
            "original_texts": [
                str(item.get("original_text", "") or "") for item in group
            ],
            "issues": [str(item.get("issue", "") or "") for item in group],
            "suggestions": [str(item.get("suggestion", "") or "") for item in group],
        }
        user_content = (
            MERGER_USER_PROMPT.format(
                payload=json.dumps(payload, ensure_ascii=False, indent=2)
            )
            + OUTPUT_EXAMPLE
        )
        msgs = self.build_messages(user_content)
        resp = self.run_with_loop(self.chat_async(msgs))
        data = self.parse_first_json(resp)
        issue = str(data.get("issue", "") or "").strip()
        suggestion = str(data.get("suggestion", "") or "").strip()
        if not issue and not suggestion:
            raise ValueError("empty issue/suggestion returned by model")
        return {"issue": issue, "suggestion": suggestion}


if __name__ == "__main__":
    tool = SegmentMergerTool()
    sample = [
        {
            "rule_title": "预付款审查",
            "segment_id": 0,
            "original_text": "丙方提交的下列单据经甲方、乙方审核无误后 20 个工作日内，支付该批次设备合同价格 70% 的到款项，合同生效后，丙方提供下列材料，甲方、乙方审核无误后20个工作日内支付给丙方本合同总价的10%的款项（￥458,000元，人民币大写肆拾伍万捌仟元整）作为预付款",
            "issue": "根据审查规则，预付款比例应大于等于合同总价款的20%，或约定发货前付清全款。合同第3.2.1条约定预付款比例为合同总价的10%（458,000元），未达到公司规定的20%最低比例要求，且未约定发货前付清全款，存在资金占用风险。",
            "risk_level": "H",
            "suggestion": "将预付款比例修改为合同总价的20%。建议修改为：“合同生效后，丙方提供下列材料，甲方、乙方审核无误后20个工作日内支付给丙方本合同总价的20%的款项（￥916,000元，人民币大写玖拾壹万陆仟元整）作为预付款”。",
            "result": "不合格",
        },
        {
            "rule_title": "付款时间审查",
            "segment_id": 0,
            "original_text": "丙方按合同约定和交货通知单的要求交付合同设备后，现场经清点无误并验收合格，丙方提交的下列单据经甲方、乙方审核无误后 20 个工作日内，支付该批次设备合同价格 70% 的到款项",
            "issue": "根据审查规则，发货后的付款项（如到货款）若未约定发货前全额付款，必须以“到货 XX 天/月”或相似表述作为充分条件之一，且若有多个条件需提及“先到为准”。当前条款约定支付“到货款”的条件仅为“现场经清点无误并验收合格”，属于以验收结果作为唯一触发条件，未设置“到货 XX 天”的闭口时间限制。若买方拖延验收，将导致卖方收款时间无限期延后，不符合规则要求。",
            "risk_level": "H",
            "suggestion": "修改为：丙方按合同约定和交货通知单的要求交付合同设备后，现场经清点无误并验收合格，或自货物到达现场之日起【30】日内（以先到者为准），丙方提交的下列单据经甲方、乙方审核无误后 20 个工作日内，支付该批次设备合同价格 70% 的到款项。",
            "result": "不合格",
        },
        {
            "rule_title": "主体审查",
            "segment_id": 0,
            "original_text": "3.2.2 ...支付该批次设备合同价格70%的到款项； ... B．乙方开具给甲方、丙方开具给乙方的金额为该批次合同价格100%的增值税专用发票；",
            "issue": "第3.2.2条约定到货款支付比例为70%，但条款要求开具金额为该批次合同价格100%的增值税专用发票。此时发票比例（100%）高于付款比例（70%）。根据审查规则，此类情况仅在合同中明确表述“货到”或“发货完成”时才为合格。虽然3.2.2条提到“交付合同设备后...验收合格”，但开票义务在3.3.1条中表述为无条件义务，未将“货到/发货完成”作为开具全额发票的强制前置条件，导致卖方在仅收到70%款项时需开具100%发票，不符合卖方利益。",
            "risk_level": "H",
            "suggestion": '明确开票节点与付款节点的对应关系，修改为："丙方应在合同设备发货后、买方支付到货款前，向乙方开具该批次设备金额100%的增值税专用发票。"',
            "result": "不合格",
        },
    ]
    print(json.dumps(tool.run(sample), ensure_ascii=False, indent=2))
