from __future__ import annotations

import logging
import time
from .grounded_rule_ds._nodes._tupletable import NameToObject
from typing import TYPE_CHECKING

from kele.syntax import _QuestionRule

from .grounded_rule_ds import GroundedProcess, GroundedRule, GroundedRuleDS
from kele.config import GrounderConfig

if TYPE_CHECKING:
    from collections.abc import Sequence

    from kele.control.grounding_selector import GroundingFlatTermWithWildCardSelector, GroundingRuleSelector
    from kele.syntax import GROUNDED_TYPE_FOR_UNIFICATION, Question, Rule

logger = logging.getLogger(__name__)


class Grounder:
    """Main structure for the grounding process."""
    def __init__(self,
                 *,
                 rule_selector: GroundingRuleSelector,
                 term_selector: GroundingFlatTermWithWildCardSelector,
                 grounded_structure: GroundedRuleDS,
                 args: GrounderConfig | None = None):  # 这些参数是不是也进selector即可
        """
        :param fact_base: Fact base used by selectors and grounding.
        :param rule_base: Rule base used during grounding.
        :param args: Grounder configuration for selection limits and heuristics.
        """
        self.grounded_structure = grounded_structure

        self.args = args or GrounderConfig()
        # TODO: With selectors in place, enforce grounding_rules_per_step and grounding_facts_per_rule via selectors.

        self._rule_selector = rule_selector
        self._term_selector = term_selector
        self._last_selected_rules: Sequence[Rule] = []

    def __select_abstract_rule(self, question: Question, m: int = 5) -> Sequence[Rule]:
        """Select abstract rules for grounding."""
        abstract_rules: Sequence[Rule] = self._rule_selector.next_rules()
        if logger.isEnabledFor(logging.DEBUG):
            active_terms_count = self._term_selector.active_terms_count()
            logger.debug(
                "Selected %s abstract rules via selector with %s active facts.",
                len(abstract_rules),
                active_terms_count if active_terms_count is not None else "unknown",
            )
        return abstract_rules

    def __select_facts_for_abstract_rules(self, question: Question, abstract_rules: Sequence[Rule]) \
            -> Sequence[tuple[Rule, list[GROUNDED_TYPE_FOR_UNIFICATION]]]:
        """Select facts/terms used to instantiate abstract rules."""
        rule_terms_pairs: list[tuple[Rule, list[GROUNDED_TYPE_FOR_UNIFICATION]]] = [(r, self._term_selector.next_terms(r)) for r in abstract_rules]
        self._last_selected_rules = [rule for rule, _ in rule_terms_pairs]
        # 变量名字我先不改，类型先写成GROUNDED_TYPE
        if logger.isEnabledFor(logging.DEBUG):
            active_rules_count = self._rule_selector.active_rules_count()
            logger.debug(
                "Selected %s rule-term pairs via selectors with %s active rules.",
                len(rule_terms_pairs),
                active_rules_count if active_rules_count is not None else "unknown",
            )
        return rule_terms_pairs

    def _select_rule_terms_pair(self, question: Question) -> Sequence[tuple[Rule, list[GROUNDED_TYPE_FOR_UNIFICATION]]]:
        """Select rules and prepare their visible fact sets."""
        return self.__select_facts_for_abstract_rules(question=question, abstract_rules=self.__select_abstract_rule(question))

    def select_facts_rules_pair(self, question: Question) -> list[tuple[Rule, list[GROUNDED_TYPE_FOR_UNIFICATION]]]:
        """Select facts first, then rules (not yet enabled)."""
        raise NotImplementedError

    def _select_grounded_rule(self, grounded_rules: Sequence[GroundedRule]) -> Sequence[GroundedRule]:
        """Filter rules that have no grounding instances (future refinement hook)."""
        self.max_grounded_series: int  # 如果注释中的TODO完成，那以后会有一个参数用于控制筛选算法或数量
        return list(grounded_rules)

    def grounding_process(self, question: Question) -> Sequence[GroundedRule]:
        """
        Run a full grounding cycle and return grounded rules.

        :param question: Question and its premises.
        :returns: Grounded rules generated for execution.
        """
        logger.info("Starting grounding process for question: %s", question.description)
        start_time = time.time()

        selected_rule_terms_pair = self._select_rule_terms_pair(question)
        with GroundedProcess(grounded_structure=self.grounded_structure, cur_rules_terms=selected_rule_terms_pair) as grounded_structure:
            grounded_structure.exec_grounding()

        grounded_rules = self.grounded_structure.get_corresponding_grounded_rules(
            abstract_rules=[r[0] for r in selected_rule_terms_pair])

        elapsed = time.time() - start_time
        logger.info("Grounding completed in %.2f seconds. Grounded rules: %s", elapsed, len(grounded_rules))

        return self._select_grounded_rule(grounded_rules)

    def reset(self) -> None:
        """Reset grounder state for a new question."""
        self.grounded_structure.reset()  # 现在的RuleCheckGraph是每条GroundedRule自己维护的，所以清空Pool就相当于删除了Graph
        # 如果后期Graph合并时，这里应当进行额外的reset过程
        self._rule_selector.reset()  # 不会清空规则，只是重置到初始状态
        # 1. 由更换问题导致的rule base整个变更，应当重新声明InferenceEngine类
        # 2. 由rule base或其他地方选择推理所用规则导致的变更，由各处自行通过set rules进行变更
        self._term_selector.reset()  # 会清空term候选表，因为事实是从初始阶段逐步推理而得的
        NameToObject.reset()

    def selected_only_question_rules(self) -> bool:
        """Return whether the last selection contained only question rules."""
        return bool(self._last_selected_rules) and all(
            isinstance(rule, _QuestionRule) for rule in self._last_selected_rules
        )
