from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import TYPE_CHECKING

from kele.control.run_mode import RunMode
from ._rule_strategies import get_rule_strategy_class

if TYPE_CHECKING:
    from collections.abc import Sequence

    from kele.syntax import Rule, _QuestionRule

    from ._rule_strategies.strategy_protocol import Feedback, RuleSelectionStrategy


@dataclass
class _RuleSelectorProfile:
    strategy: RuleSelectionStrategy
    rules: list[Rule]
    # TODO: question-specific state should move into a dedicated question mode profile.
    question_rules: list[Rule]
    # TODO: question-specific state should move into a dedicated question mode profile.
    question_rule_interval: int
    used_rule_cnt: int
    at_fixpoint: bool


class GroundingRuleSelector:
    """
    对外统一入口，内部按 mode 委托给策略实现。
    - 允许切换策略（接口不限制“连续取中”）
    - 允许更新规则集合
    """

    def __init__(self, strategy: str = "sequential_cyclic", question_rule_interval: int = -1) -> None:
        self._mode: RunMode = RunMode.INFERENCE
        self._strategy_name_by_mode: dict[RunMode, str] = dict.fromkeys(RunMode, strategy)
        self._profiles: dict[RunMode, _RuleSelectorProfile] = {
            mode: self._build_profile(strategy_name=strategy)
            for mode in RunMode
        }
        # question_rule_interval 当前只作用于 INFERENCE mode。
        self._profiles[RunMode.INFERENCE].question_rule_interval = question_rule_interval

    @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) -> _RuleSelectorProfile:
        strategy_cls = get_rule_strategy_class(strategy_name)
        return _RuleSelectorProfile(
            strategy=strategy_cls(),
            rules=[],
            question_rules=[],
            question_rule_interval=-1,
            used_rule_cnt=0,
            at_fixpoint=False,
        )

    def set_at_fixpoint(self, *, at_fixpoint: bool) -> None:
        """设置当前 mode profile 是否处于不动点。"""
        self._profiles[self.mode].at_fixpoint = at_fixpoint

    def next_rules(self) -> Sequence[Rule]:
        """按当前 mode 选择本轮待 grounding 的规则。

        :raises ValueError: 当 question_rule_interval 小于 1 且不为 -1 时抛出。
        """
        mode = self.mode
        profile = self._profiles[mode]

        if mode != RunMode.INFERENCE:
            if not profile.rules:
                return []
            return profile.strategy.select_next()

        # TODO: 逐步将 fixpoint 处理迁移到独立的 question mode，
        # 避免与 INFERENCE mode 耦合。
        if profile.at_fixpoint:
            if profile.question_rules:
                profile.used_rule_cnt = 0
            return profile.question_rules

        # 确定实际的检查间隔（仅 inference mode）
        interval = profile.question_rule_interval
        if interval < 1 and interval != -1:
            raise ValueError(
                "question_rule_interval must be >= 1, or -1 to use the total count of normal rules as the interval",
            )

        if interval == -1:
            normal_count = len(profile.rules)
            interval = normal_count or 1

        if profile.used_rule_cnt >= interval and profile.question_rules:
            profile.used_rule_cnt = 0
            return profile.question_rules

        if not profile.rules:
            return []
        rules = profile.strategy.select_next()
        profile.used_rule_cnt += len(rules)
        return rules

    def active_rules_count(self) -> int:
        """返回当前 mode 下候选规则数量。"""
        return len(self._profiles[self.mode].rules)

    def mode_rules_count(self, mode: RunMode) -> int:
        """返回指定 mode 下候选规则数量。"""
        return len(self._profiles[mode].rules)

    def reset(self) -> None:
        """重置所有 mode 的策略与规则缓存。"""
        self.mode = RunMode.INFERENCE
        for profile in self._profiles.values():
            profile.strategy.reset()
            profile.rules = []
            profile.question_rules = []
            profile.used_rule_cnt = 0
            profile.at_fixpoint = False

    def set_strategy(self, strategy: RuleSelectionStrategy) -> None:
        """替换当前 mode 的策略实现。

        若当前 mode 已有规则，会调用 ``set_rules`` 同步规则；
        是否重置策略内部状态由具体策略实现决定。
        """
        profile = self._profiles[self.mode]
        profile.strategy = strategy
        if profile.rules:
            profile.strategy.set_rules(profile.rules)
        else:
            warnings.warn("No given rules, please call set_rules after calling set_strategy.", stacklevel=2)

    def set_mode_strategy(self, mode: RunMode, strategy_name: str) -> None:
        """替换指定 mode 的策略实现。"""
        self._strategy_name_by_mode[mode] = strategy_name
        profile = self._profiles[mode]
        strategy_cls = get_rule_strategy_class(strategy_name)
        profile.strategy = strategy_cls()
        if profile.rules:
            profile.strategy.set_rules(profile.rules)

    def set_mode_rules(
        self,
        *,
        mode: RunMode,
        rules: Sequence[Rule],
        # TODO: question_rules should be owned by a dedicated question mode profile.
        question_rules: Sequence[_QuestionRule] | None = None,
        # TODO: question_rule_interval should be owned by a dedicated question mode profile.
        question_rule_interval: int | None = None,
    ) -> None:
        """更新指定 mode 的规则集合并同步策略状态。"""
        profile = self._profiles[mode]
        profile.rules = list(rules)
        if question_rules is not None:
            profile.question_rules = list(question_rules)
        if question_rule_interval is not None:
            profile.question_rule_interval = question_rule_interval

        # 重置计数器
        profile.used_rule_cnt = 0
        profile.at_fixpoint = False

        profile.strategy.reset()
        if profile.rules:
            profile.strategy.set_rules(profile.rules)

    def set_rules(
        self,
        normal_rules: Sequence[Rule],
        # TODO: keep for backward compatibility; move question rules to dedicated mode setup later.
        question_rules: Sequence[_QuestionRule],
        *,
        conflict_rules: Sequence[Rule] | None = None,
    ) -> None:
        """一次性设置 inference/conflict 两类规则。"""
        self.set_mode_rules(
            mode=RunMode.INFERENCE,
            rules=normal_rules,
            question_rules=question_rules,
        )
        self.set_mode_rules(
            mode=RunMode.CONFLICT_CHECK,
            rules=conflict_rules or [],
            question_rules=[],
            question_rule_interval=-1,
        )

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