from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, cast, overload

from kele.knowledge_bases.builtin_base.builtin_facts import true_const
from kele.syntax import Constant, Variable
from kele.syntax.base_classes import Assertion, CompoundTerm, Formula, Question, Rule, _QuestionRule
from kele.syntax.dnf_converter import RuleSafetyProcesser

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_suffix: int = 1

        self._var_counter = 0

        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.

        Clears active rules and marks the base as uninitialized.
        """
        self.cur_rules.clear()
        self.initial_sign = False

    def _preprocess_rule[T1: Rule | _QuestionRule](self, rule: T1) -> Sequence[T1]:
        return [self._resolve_variable_collisions(single_rule) for single_rule in
                self._split_and_rename_rule(rule)]

    def _resolve_variable_collisions[T1: Rule | _QuestionRule](self, rule: T1) -> T1:
        """Rename variables in a rule to avoid collisions."""
        var_map: dict[str, Variable] = {}  # 仅作为临时工作区
        new_head = self._generate_renamed_item(rule.head, var_map)
        new_body = self._generate_renamed_item(rule.body, var_map)

        return cast("T1", rule.replace(head=new_head, body=new_body))

    def _split_and_rename_rule[T1: Rule](self, rule: T1) -> Sequence[T1]:
        """Assign names to unnamed rules and split rules into DNF."""
        origin_name = rule.name or f"rule_{self.rule_suffix}"
        self.rule_suffix += 1

        rule_spliter = RuleSafetyProcesser()  # TODO: 这里每次创建一个TODO还是有点奇怪
        splited_rules = rule_spliter.split_rule_and_process_safety(rule)

        if len(splited_rules) == 1:
            return [splited_rules[0].replace(name=origin_name)]

        result = []
        for idx, single_rule in enumerate(splited_rules):
            new_name = f"{origin_name}_{idx + 1}"
            result.append(single_rule.replace(name=new_name))

        return result

    @overload
    def _generate_renamed_item(self, item: Variable, var_map: dict[str, Variable]) -> Variable: ...
    @overload
    def _generate_renamed_item(self, item: Constant, var_map: dict[str, Variable]) -> Constant: ...
    @overload
    def _generate_renamed_item(self, item: CompoundTerm, var_map: dict[str, Variable]) -> CompoundTerm: ...
    @overload
    def _generate_renamed_item(self, item: Assertion, var_map: dict[str, Variable]) -> Assertion: ...
    @overload
    def _generate_renamed_item(self, item: Formula, var_map: dict[str, Variable]) -> Formula: ...
    @overload
    def _generate_renamed_item(self, item: None, var_map: dict[str, Variable]) -> None: ...

    def _generate_renamed_item(self,
                               item: Formula | Assertion | CompoundTerm | Variable | Constant | None,
                               var_map: dict[str, Variable]) -> (
            Formula | Assertion | CompoundTerm | Variable | Constant | None):
        if isinstance(item, Variable):
            if item.symbol not in var_map:
                new_name = f"_v{self._var_counter}"
                self._var_counter += 1
                var_map[item.symbol] = item.create_renamed_variable(new_name)
                # risk: 两条规则里如果都用到了x，然后使用者使用了同一个x的instance，
                # 我们是否有理由认为它俩是一个x？（我认为不合理，所以我们应当不管instance，而只是取它的value本身。但这样会不会给使用者带来困惑？
                # 他可能刻意用了同一个instance，甚至在里面存了一些自定义的信息。但我们可能是直接丢弃然后换个了新名字新地址）
            return var_map[item.symbol]

        if isinstance(item, Constant):
            return item

        if isinstance(item, CompoundTerm):
            return CompoundTerm.from_parts(
                item.operator,
                tuple(self._generate_renamed_item(arg, var_map) for arg in item.arguments),
            )

        if isinstance(item, Assertion):
            return Assertion.from_parts(
                self._generate_renamed_item(item.lhs, var_map),
                self._generate_renamed_item(item.rhs, var_map),
            )

        if isinstance(item, Formula):
            return Formula(
                self._generate_renamed_item(item.formula_left, var_map),
                item.connective,
                self._generate_renamed_item(item.formula_right, var_map),
            )

        if item is None:
            return None

        raise ValueError(f"Unknown item type: {type(item)}")

    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}]"
