from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

from kele.syntax.base_classes import Assertion

if TYPE_CHECKING:
    from collections.abc import Sequence

    from kele.config import KBConfig
    from kele.control.grounding_selector import GroundingFlatTermWithWildCardSelector
    from kele.equality import Equivalence
    from kele.syntax.base_classes import FACT_TYPE, Question
    from kele.syntax.external import SankuManagementSystem


# FIXME: 存入事实库和规则库的信息应当满足safe的条件，我们将在 #132 中对CompoundTerm加入generic且引入safe的约束，以完成更精细的类型标注
class FactBase:
    """Store and manage facts for the inference engine.

    :param initial_facts_or_dir_or_path: Initial facts or path to load them from.
    :param equivalence_handler: Equivalence manager for equality updates.
    :param term_selector: Selector used for grounding term candidates.
    :param sk_system_handler: External knowledge system handler.
    :param args: Knowledge-base configuration.
    """

    def __init__(self,
                 initial_facts_or_dir_or_path: Sequence[FACT_TYPE] | str,
                 equivalence_handler: Equivalence,
                 term_selector: GroundingFlatTermWithWildCardSelector,
                 sk_system_handler: SankuManagementSystem,
                 args: KBConfig,
                 ):
        if isinstance(initial_facts_or_dir_or_path, str):
            self.facts: set[FACT_TYPE] = self._read_facts(initial_facts_or_dir_or_path)
        else:
            self.facts = set(initial_facts_or_dir_or_path)

        for fact in self.facts:
            self._validate_fact_for_storage(fact, check_free_variables=True)

        self.equivalence_handler = equivalence_handler
        self.term_selector = term_selector
        self.sk_system_handler = sk_system_handler  # HACK: 三库系统在解题前，会执行一次initial_by_question更新事实库。考虑到
        # 事实库的初始化是一个独立环节，此函数的调用被放置到FactBase内，而非最外层的InferenceEngine内。虽然从模块拆分的角度上FactBase和
        # 三库系统是两个系统，可能会建议在main函数内完成二者的交互。
        self._args = args
        self.max_facts = self._args.fact_cache_size

        self.initial_sign = False  # 引擎希望先挑选FactBase中的部分事实进行初始化。因此在__init__阶段我们认为初始化未完成，为False。
        # 当初始化完成后、当前值为True时才应当进行后续其他操作
        self.cur_facts: set[FACT_TYPE] = set()

    def _read_facts(self, path: str) -> set[FACT_TYPE]:
        """Load facts from a directory or file.

        :param path: Directory or file path.
        :returns: Set of facts.
        :raise NotImplementedError: will be completed before 0.0.1
        """
        raise NotImplementedError

    def _add_or_not(self, fact: FACT_TYPE) -> bool:
        """Decide whether a fact should be added to the fact base."""

        # 1. 判断事实是否已存在
        if fact in self.facts:
            return False  # 已存在的事实不再加入

        # 2. 判断是否超过事实库大小上限
        if self.max_facts != -1 and len(self.facts) >= self.max_facts:  # noqa: SIM103  # 函数还可以继续扩充，
            # 不应当直接简化为return (self.max_facts is not None...)，因此注释掉SIM103
            return False  # TODO: 如果事实库已满，最好是根据启发式决定是否丢弃现有事实或丢弃库内事实但保留当前事实

        # 3. 不合需求没必要保留（依赖启发式/NN判断）

        return True

    def add_facts(
        self,
        facts: Sequence[FACT_TYPE],
        *,
        force_add: bool = False,
        check_free_variables: bool = False,
    ) -> list[FACT_TYPE]:
        """
        Add facts and update equivalence classes.

        :param facts: Facts to add.
        :param force_add: Force insertion even if heuristics reject it.
        :param check_free_variables: Validate that facts have no free variables.
        :returns: Facts that were actually added.
        """
        added_facts = []
        for fact in facts:
            self._validate_fact_for_storage(fact, check_free_variables=check_free_variables)

        if force_add:  # 冗余一点避免多次判断force_add
            for fact in facts:
                self.facts.add(fact)
                added_facts.append(fact)
        else:
            for fact in facts:
                if self._add_or_not(fact):
                    self.facts.add(fact)
                    added_facts.append(fact)

        self.cur_facts |= set(added_facts)
        self.equivalence_handler.update_equiv_class(added_facts)
        if added_facts:
            self.term_selector.update_terms(facts=added_facts)

        return added_facts

    @staticmethod
    def _validate_fact_for_storage(fact: FACT_TYPE, *, check_free_variables: bool) -> None:
        if not isinstance(fact, Assertion):
            raise TypeError(f"Fact {fact} is not an Assertion, which is not allowed in the fact base.")
        if check_free_variables and fact.free_variables:
            raise ValueError(f"Fact {fact} contains free variables, which is not allowed.")

    def _select_initial_facts(self, question: Question, topn: int | None = None) -> list[FACT_TYPE]:
        """Select the most relevant facts for a question.

        :param question: Question being solved.
        :param topn: Limit on the number of facts to return; ``None`` disables the limit.
        :returns: Selected facts.
        """
        all_facts = list(self.facts)

        if topn is None or topn == -1 or topn >= len(all_facts):
            return all_facts

        return all_facts[:topn]  # TODO: 默认按插入顺序（或原始顺序）截取，是优化点

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

        :param question: Question being solved.
        :param topn: Optional cap on facts selected from the base.
        """
        selected_facts = self._select_initial_facts(question, topn)
        self.equivalence_handler.update_equiv_class(selected_facts)
        self.cur_facts |= set(selected_facts)

        sk_facts = self.sk_system_handler.initial_by_question(question, topn)
        self.add_facts(sk_facts)

        self.cur_facts |= set(sk_facts)
        self.initial_sign = True

    def reset_fact_base(self) -> None:
        """Reset the fact base to its pre-initialized state."""
        self.initial_sign = False
        self.cur_facts.clear()
        self.equivalence_handler.clear()  # 理论上如果这个函数仅有main调用，这一行是不应当有的。但我担心由于函数是外部的，如果被
        # 凑巧谁的代码调用过时，一步clear的冗余可以减少风险

    def get_facts(self) -> list[FACT_TYPE]:
        """Return facts used or derived up to the current round."""
        if self.initial_sign:
            return list(self.cur_facts)

        warnings.warn("Fact base has not been initialized yet.", stacklevel=2)
        return list(self.facts)

    def __str__(self) -> str:
        fact_count = len(self.facts)
        show_topn = 5

        fact_summary = ', '.join(str(fact) for fact in list(self.facts)[:show_topn])  # 只展示前五条事实

        # 如果事实数目大于5，提示用户后续还有更多
        if fact_count > show_topn:
            fact_summary += "..."

        return f"FactBase with {fact_count} facts. First 5 facts: [{fact_summary}]"
