from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

from kele.knowledge_bases.builtin_base.builtin_facts import true_const
from kele.knowledge_bases.rule_utils import RulePreprocessor
from kele.syntax.base_classes import Assertion, Question, Rule, _QuestionRule

if TYPE_CHECKING:
    from collections.abc import Sequence

    from kele.config import KBConfig


class RuleBase:
    """Store and manage abstract rules for inference.

    :param rules_or_dir_or_path: Rules or path to load them from.
    :param args: Knowledge-base configuration.
    """

    def __init__(self, rules_or_dir_or_path: Sequence[Rule] | str, args: KBConfig):
        if isinstance(rules_or_dir_or_path, str):
            self.rules = list(dict.fromkeys(self._read_rules(rules_or_dir_or_path)))
            # 使用list代替set，以保持输入的规则顺序、因为set后的随机顺序会引入
            # 不必要的不确定性，随意的排序也对于推理无益
        else:
            self.rules = list(dict.fromkeys(rules_or_dir_or_path))  # 加入去重操作

        self._rule_preprocessor = RulePreprocessor()

        new_rules: list[Rule] = []
        for single_rule in self.rules:
            new_rules.extend(self._preprocess_rule(single_rule))
        self.rules = new_rules

        self.initial_sign: bool = False  # initial_rule_base后为True，和Fact base一样
        self.cur_rules: list[Rule] = []  # FactBase由于经常需要更新而用了set，而abstract rule不变，为了便捷用了list
        self.question_rule: list[_QuestionRule] = []  # 用于存储question转化成的rule，不干涉正常流程，不参与select等。

        self._args = args

    def _read_rules(self, path: str) -> list[Rule]:
        """Load rules from a directory or file.

        :param path: Directory or file path.
        :returns: Loaded rules.
        :raise NotImplementedError: will be completed before 0.0.1
        """
        raise NotImplementedError

    def add_rule(self, rule: Rule) -> None:
        """Add a rule at runtime (currently not supported)."""
        warnings.warn(
            "Not supported yet. When runtime rule additions are enabled, SelectionStrategy must add add_rule and "
            "adjust RuleCheckRule behavior as needed.",
            stacklevel=2,
        )
        if rule not in self.rules:
            self.rules.extend(self._preprocess_rule(rule))

    def get_rules(self) -> list[Rule]:
        """Return the active rules for the current run."""
        if self.initial_sign:
            return list(self.cur_rules)

        warnings.warn("Rule base has not been initialized yet.", stacklevel=2)
        return [r for r in self.rules if not isinstance(r, _QuestionRule)]

    def get_question_rules(self) -> list[_QuestionRule]:
        """Return rules derived from the current question."""
        return self.question_rule

    def _select_rules_from_current_question(self, question: Question, topn: int | None = None) -> list[Rule]:
        """
        Select rules based on the current question.

        :param question: Current question.
        :param topn: Optional limit on the number of rules to return.
        :returns: Selected rules.
        """
        # 由于选前k个而需要list或tuple一下  TODO: 另外：暂时将所有规则视为可能相关的规则，以后可以添加过滤逻辑或排序什么的。
        all_rules = list(self.rules)

        # 如果没有指定限制则返回所有规则。
        if topn is None or topn == -1 or topn >= len(all_rules):
            return all_rules

        return all_rules[:topn]

    def initial_rule_base(self, question: Question, topn: int | None = None) -> None:
        """Initialize the active rule set for a new question.

        :param question: Question being solved.
        :param topn: Optional cap on the number of selected rules.
        """
        if question.question:
            question_r = _QuestionRule.from_parts(
                head=Assertion.from_parts(_QuestionRule.QUESTION_SOLVED_FLAG, true_const),
                body=question.question,
                name=_QuestionRule.QUESTIONRULE_NAME,
            )
            self.question_rule.extend(self._preprocess_rule(question_r))

        # TODO: consider moving question rules handling to main.py to avoid mixing them into selector/base flow.
        selected_rules = self._select_rules_from_current_question(question, topn) + self.question_rule
        self.cur_rules = [rule for r in selected_rules for rule in self._preprocess_rule(r)]
        self.initial_sign = True

    def reset_rule_base(self) -> None:
        """Reset to the pre-initialized state."""
        self.cur_rules.clear()
        self.initial_sign = False

    def _preprocess_rule[T1: Rule](self, rule: T1) -> list[T1]:
        # 预处理逻辑已抽离到 RulePreprocessor，避免在 RuleBase 内重复维护。
        return self._rule_preprocessor.preprocess_rule(rule)

    def __str__(self) -> str:
        rule_count = len(self.rules)

        show_topn = 5

        # 展示前5条规则
        preview_rules = ', '.join(str(rule) for rule in list(self.rules)[:5])
        if rule_count > show_topn:
            preview_rules += "..."

        return f"RuleBase with {rule_count} rules. First 5 rules: [{preview_rules}]"
