from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Any
import re

from core.config import ALL_RULESET_IDS, DEFAULT_RULESET_ID
from core.tool import ToolBase, tool, tool_func
from utils.excel_util import ExcelUtil


@tool("retrieve_reference", "审查参考检索")
class RetrieveReferenceTool(ToolBase):
    def __init__(self) -> None:
        self.default_ruleset_id = DEFAULT_RULESET_ID
        self.column_map = {
            "id": "ID",
            "title": "审查项",
            "rule": "审查规则",
            "level": "风险等级",
            "triggers": "触发词",
            "suggestion_template": "建议模板",
            "case": "案例",
            "summary": "摘要项",
        }
        rules_path = (
            Path(__file__).resolve().parent.parent.parent / "data" / "rules.xlsx"
        )
        self.rulesets: Dict[str, List[Dict[str, Any]]] = {}
        for rs_id in ALL_RULESET_IDS:
            rules = ExcelUtil.load_mapped_excel(
                rules_path, sheet_name=rs_id, column_map=self.column_map
            )
            self.rulesets[rs_id] = rules

    @tool_func(
        {
            "type": "object",
            "properties": {
                "ruleset_id": {"type": "string"},
                "routed_rule_titles": {"type": "array", "items": {"type": "string"}},
            },
            "required": [],
        }
    )
    def run(
        self, ruleset_id: str = "", routed_rule_titles: List[str] | None = None
    ) -> Dict[str, Any]:
        target_ruleset_id = ruleset_id or self.default_ruleset_id
        full_rules = (
            self.rulesets.get(target_ruleset_id)
            or self.rulesets.get(self.default_ruleset_id, [])
            or []
        )
        if routed_rule_titles is None:
            rules = full_rules
        else:
            title_set = {
                title for title in routed_rule_titles if isinstance(title, str)
            }
            rules = [r for r in full_rules if r.get("title") in title_set]

        return {
            "ruleset_id": target_ruleset_id,
            "rules": rules,
            "rule_titles": [r.get("title", "") for r in rules],
            "total": len(rules),
        }

    def summary_keywords(self, rules: List[Dict[str, Any]]) -> List[str]:
        return [r.get("summary", "") for r in rules if r.get("summary")]


if __name__ == "__main__":
    tool = RetrieveReferenceTool()
    result = tool.run(ruleset_id="金盘", routed_rule_titles=None)
    for rule in result.get("rules", []):
        print(f"Rule Title: {rule.get('title')}")
        print(f"Case: {rule.get('case')}")
        print("-" * 20)
    # print(result.get("total", 0))
