from __future__ import annotations

import itertools
from dataclasses import dataclass
from typing import TYPE_CHECKING

from kele.control.grounding_selector._selector_utils import _unify_ground_terms_from_rules, _unify_into_terms
from kele.control.run_mode import RunMode

from ._term_strategies import get_term_strategy_class

if TYPE_CHECKING:
    from collections.abc import Sequence

    from kele.config import Config
    from kele.equality import Equivalence
    from kele.syntax import FACT_TYPE, GROUNDED_TYPE_FOR_UNIFICATION, Rule

    from ._term_strategies.strategy_protocol import Feedback, TermSelectionStrategy


@dataclass
class _TermSelectorProfile:
    strategy_name: str
    strategy: TermSelectionStrategy


class GroundingFlatTermWithWildCardSelector:
    # 此时是 FlatTerm-level grounding，所以进 strategy 不止是 Term。
    # 但 strategy 本身是 Term 和 FlatTerm 均合适的，所以标准仍然是 TERM_TYPE。
    # 此外，目前通配符和 flat term 无法分割，所以命名强调了 wild card。
    """对外统一入口，内部委托给策略实现。"""

    def __init__(self, equivalence: Equivalence, args: Config) -> None:
        self._mode: RunMode = RunMode.INFERENCE
        self._equivalence = equivalence
        self._args = args
        # 全局 term 快照（去重集合）:
        # - 常规路径下，update_terms* 会把“新增 term”直接同步到所有 profile.strategy；
        # - 仅当某个 mode 的 strategy 被重建（set_mode_strategy）时，需要恢复完整 term 集。
        # 因此这里不是“每轮都全量 add”，而是“重置后的状态恢复来源”。
        self._all_terms: set[GROUNDED_TYPE_FOR_UNIFICATION] = set()
        strategy_name = args.strategy.grounding_term_strategy
        self._profiles: dict[RunMode, _TermSelectorProfile] = {
            mode: self._build_profile(strategy_name) for mode in RunMode
        }

    @property
    def mode(self) -> RunMode:
        """当前运行模式。"""
        return self._mode

    @mode.setter
    def mode(self, mode: RunMode) -> None:
        if not isinstance(mode, RunMode):
            raise TypeError(f"mode must be RunMode, got {type(mode).__name__}")
        self._mode = mode

    def set_mode(self, mode: RunMode) -> None:
        """显式设置当前运行模式（等价于 ``self.mode = mode``）。"""
        self.mode = mode

    @staticmethod
    def _build_profile(strategy_name: str) -> _TermSelectorProfile:
        strategy_cls = get_term_strategy_class(strategy_name)
        return _TermSelectorProfile(strategy_name=strategy_name, strategy=strategy_cls())

    def set_mode_strategy(self, mode: RunMode, strategy_name: str) -> None:
        """设置指定 mode 的 term 选择策略。

        说明:
        - 当前主流程通常不主动切换 strategy；该接口主要用于扩展/实验配置。
        - 新 strategy 实例初始为空，需要回放已有 term 快照，避免切换后候选集丢失。
        """
        profile = self._build_profile(strategy_name)
        if self._all_terms:
            # 新建 profile 时需要恢复全量快照；否则只保留切换后的增量 term。
            profile.strategy.add_terms(list(self._all_terms))
        self._profiles[mode] = profile

    def _record_terms(self, terms: Sequence[GROUNDED_TYPE_FOR_UNIFICATION]) -> None:
        """更新全局 term 快照（集合去重累计）。"""
        self._all_terms.update(terms)

    def next_terms(
        self,
        rule: Rule,
    ) -> list[GROUNDED_TYPE_FOR_UNIFICATION]:
        """为给定规则选择候选事实/term 用于 grounding。"""
        mode = self.mode
        profile = self._profiles[mode]
        init_terms = list(profile.strategy.select_next(rule))
        if self._args.run.semi_eval_with_equality:
            selected_terms: set[GROUNDED_TYPE_FOR_UNIFICATION] = set()
            for term in init_terms:
                if term not in selected_terms:
                    equiv_terms = self._equivalence.get_related_item(term)
                    selected_terms |= set(equiv_terms)  # TODO: 这里后面考虑是否把_unify_into_flat_terms挪到term selector

                    # 在semi-evaluation中纳入等价类其余terms，这是由
                    # 断言逻辑带来的内嵌、强制等词公理，和term-level grounding二者共同决定的算法，因为不放置于具体的Strategy中，而是作为
                    # Selector自身的行为。
            return list(selected_terms)

        return list(set(init_terms))

    def active_terms_count(self) -> int | None:
        """返回当前候选 term 数量，若策略不支持则返回 None。"""
        # TODO: remove direct strategy internals access (_terms) after exposing a typed strategy API.
        terms = getattr(self._profiles[self.mode].strategy, "_terms", None)
        if terms is None:
            return None
        return len(terms)

    def reset(self) -> None:
        """重置选择器（当前实现会重置所有 mode 的 term 选择状态）。"""
        self.mode = RunMode.INFERENCE
        for profile in self._profiles.values():
            profile.strategy.reset()
        self._all_terms = set()

    def update_terms(
        self,
        terms: Sequence[GROUNDED_TYPE_FOR_UNIFICATION] | None = None,
        facts: Sequence[FACT_TYPE] | None = None,
    ) -> None:
        """更新事实集合，并同步给所有 mode profile 的策略。

        常规增量路径: 这里直接把本轮新 term 推送到全部 profile.strategy；
        _all_terms 只用于 strategy 重建后的恢复，不参与日常重复同步。

        :raise ValueError: 没有可选 terms/facts 时抛出。
        """
        if terms is not None:
            fact_terms = list(terms)
        elif facts is not None:
            # 更进一步地，其实只有 Assertion 能触发等式公理。
            fact_terms = list(itertools.chain.from_iterable(_unify_into_terms(fact) for fact in facts))
        else:
            raise ValueError("terms or facts cannot be empty")
        if not fact_terms:
            return

        # 不应有 action operator 的 term。fact_terms = [f for f in fact_terms if f.operator.implement_func is not None]
        for profile in self._profiles.values():
            profile.strategy.add_terms(fact_terms)
        self._record_terms(fact_terms)

    def update_terms_from_rules(self, rules: Sequence[Rule]) -> None:
        """将规则/问题中无变量的 term 加入候选表。"""
        rule_terms = _unify_ground_terms_from_rules(rules)
        if not rule_terms:
            return
        for profile in self._profiles.values():
            profile.strategy.add_terms(rule_terms)
        self._record_terms(rule_terms)

    def send_feedback(self, feedback: Feedback) -> None:
        """把一次选择后的反馈转发给指定 mode 的当前策略。"""
        self._profiles[self.mode].strategy.on_feedback(feedback)
