from __future__ import annotations

import re
from typing import Dict, List, Optional,Tuple

from core.tool import ToolBase, tool, tool_func
from core.memory import MemoryStore

TOPIC_KEYWORDS = {
    "付款": ["付款", "支付", "价款", "结算", "发票", "对账", "款项", "账期"],
    "验收": ["验收", "确认", "签收", "测试", "交付确认", "验收单", "验收报告"],
    "交付": ["交付", "交货", "交付物", "交付时间", "里程碑", "上线", "交接"],
    "违约": ["违约", "违约金", "赔偿", "损失", "承担责任", "罚金", "扣款"],
    "解除": ["解除", "终止", "解约", "撤销"],
    "责任限制": ["责任上限", "间接损失", "免责", "赔偿上限", "责任限制", "不可抗力"],
    "保密": ["保密", "保密信息", "披露", "泄露", "保密期限", "例外"],
    "知识产权": ["知识产权", "著作权", "专利", "商标", "成果", "源代码", "许可"],
    "争议解决": ["争议", "仲裁", "诉讼", "管辖", "适用法律", "法院"],
}
DEFAULT_TOP_K = 6

@tool("memory_read", "分段记忆读取")
class MemoryReadTool(ToolBase):
    @tool_func(
        {
            "type": "object",
            "properties": {
                "segment_ids": {"type": "array"},
                "topics": {"type": "array"},
                "query": {"type": "string"},
                "top_k": {"type": "int"},
                "mode": {"type": "string"},
            },
            "required": ["top_k"],
        }
    )

    def parse_neighbor_ids(self, segment_id: int) -> List[int]:
        neighbor_ids = []
        if segment_id >= 1:
            neighbor_ids.append(segment_id - 1)
        if segment_id >= 2:
            neighbor_ids.append(segment_id - 2)
        return neighbor_ids
    
    def normalize_text(self, s: str) -> str:
        return re.sub(r"\s+", " ", s.strip())
    
    def derive_topics_rule_based(self, segment_text: str, max_topics: int = 3) -> List[str]:
        text = self.normalize_text(segment_text)
        scores: List[Tuple[str, int]] = []
        for topic, kws in TOPIC_KEYWORDS.items():
            hit = sum(1 for kw in kws if kw in text)
            if hit > 0:
                scores.append((topic, hit))
        scores.sort(key=lambda x: x[1], reverse=True)
        return [t for t, _ in scores[:max_topics]]


    def run(self, store: MemoryStore, segment_id: int, segment_text: str,
                     top_k: int = DEFAULT_TOP_K) -> Dict:
        neighbor_ids = self.parse_neighbor_ids(segment_id)
        topics = self.derive_topics_rule_based(segment_text)

        if not topics:
            topics = ["其他"]

        candidates = store.list()
        ids = set(segment_ids or [])
        if ids:
            candidates = [c for c in candidates if c.segment_id in ids]
        if topics:
            candidates = [c for c in candidates if set(c.topics).intersection(set(topics))]

        key = (query or "").strip().lower()
        scored: List[Dict] = []
        for c in candidates:
            hay = " ".join([
                c.summary.one_liner,
                " ".join(c.summary.structured),
                " ".join(c.topics),
                " ".join(c.dependencies),
                " ".join(c.evidence_quotes),
                " ".join([r.flag + " " + r.evidence + " " + r.suggestion for r in c.risks]),
            ]).lower()
            score = 0.0
            if key:
                tokens = [t for t in re.split(r"\s+", key) if t]
                score = sum(1.0 for t in tokens if t in hay) / max(1, len(tokens))
            rec = c.to_dict()
            rec["relevance"] = round(score, 2)
            rec["reason"] = "关键词匹配" if key else "主题/ID过滤"
            scored.append(rec)
        scored.sort(key=lambda x: x.get("relevance", 0.0), reverse=True)
        return {"memories": scored[:top_k]}
