from __future__ import annotations

import logging
from enum import Enum, auto
from typing import TYPE_CHECKING, Literal

from kele.config import ExecutorConfig, RunControlConfig
from kele.syntax import Variable, _QuestionRule

if TYPE_CHECKING:
    from collections.abc import Mapping, Sequence

    from kele.equality import Equivalence
    from kele.syntax import FACT_TYPE, CompoundTerm, Constant, Question, Rule, SankuManagementSystem

logger = logging.getLogger(__name__)


class InferenceStatus(Enum):
    """推理过程的终止状态类型"""

    SUCCESS = auto()  # 成功推理出答案
    MAX_STEPS_REACHED = auto()  # 达到最大执行轮次限制
    MAX_ITERATIONS_REACHED = auto()  # 达到最大迭代次数限制（从main循环）
    FIXPOINT_REACHED = auto()  # 达到不动点
    CONFLICT_DETECTED = auto()  # 推理中发现矛盾
    EXTERNALLY_INTERRUPTED = auto()  # 被外部中断（如人为终止）
    NO_MORE_RULES = auto()  # 没有更多实例化规则可以推理
    CONTINUE = auto()  # 当前轮未结束，继续执行

    def log_message(self) -> str:
        """每个状态对应的日志输出信息"""
        return {
            InferenceStatus.SUCCESS: "Successfully inferred the answer.",
            InferenceStatus.MAX_STEPS_REACHED: "Max execution steps reached, terminating.",
            InferenceStatus.MAX_ITERATIONS_REACHED: "Max inference iterations reached, terminating.",
            InferenceStatus.FIXPOINT_REACHED: "Fixpoint reached, terminating.",
            InferenceStatus.CONFLICT_DETECTED: "Conflict detected during execution.",
            InferenceStatus.EXTERNALLY_INTERRUPTED: "Execution was externally interrupted.",
            InferenceStatus.NO_MORE_RULES: "No more grounded rules available for execution.",
            InferenceStatus.CONTINUE: "Execution continues to next round."
        }[self]

    def is_terminal_for_executor(self) -> bool:
        """判断执行器是否应该终止"""
        return self != InferenceStatus.CONTINUE

    def is_terminal_for_main_loop(self) -> bool:
        """判断主循环是否应该终止"""
        return self in {
            InferenceStatus.SUCCESS,
            InferenceStatus.MAX_STEPS_REACHED,
            InferenceStatus.MAX_ITERATIONS_REACHED,
            InferenceStatus.FIXPOINT_REACHED,
            InferenceStatus.EXTERNALLY_INTERRUPTED
        }


class QuerySolutionManager:
    """查询解管理器：负责收集、扩展（等价解）、去重并按需打印查询解。"""

    def __init__(self, args: RunControlConfig) -> None:
        self.collected_solutions: list[Mapping[Variable, Constant | CompoundTerm]] = []
        self.args = args

        # 统计“已记录的唯一解”数量（包含等价扩展后的解）
        self._solution_count: int = 0

        # 记录/输出去重
        self._seen_solution_keys: set[tuple[tuple[str, str], ...]] = set()  # 去重：已出现过的解（含等价扩展）
        self._printed_solution_keys: set[tuple[tuple[str, str], ...]] = set()  # 记录已输出的解，避免重复打印
        self._pending_solutions: list[Mapping[Variable, Constant | CompoundTerm]] = []  # 待输出：新增未打印解（含等价扩展）

    def reset(self) -> None:
        """清空已收集的解、去重状态和计数。"""
        self.collected_solutions.clear()
        self._solution_count = 0
        self._seen_solution_keys.clear()
        self._printed_solution_keys.clear()
        self._pending_solutions.clear()

    # -----------------------------
    # add（仅记录，不打印）
    # -----------------------------
    def add_solution(
        self,
        solution: Mapping[Variable, Constant | CompoundTerm],
        *,
        equivalence: Equivalence | None = None,
    ) -> None:
        """记录 1 个 solution，并对其每个绑定值做等价扩展，只合入新的 solutions。

        返回：本次新增（去重后）记录的解数量。
        """
        self._record_with_equivalents(solution, equivalence=equivalence)

    def add_solutions(
        self,
        solutions: Sequence[Mapping[Variable, Constant | CompoundTerm]],
        *,
        equivalence: Equivalence | None = None,
    ) -> None:
        """批量记录 solutions（并对每个解做等价扩展）。"""
        for s in solutions:
            self.add_solution(s, equivalence=equivalence)

    def has_pending_solutions(self) -> bool:
        """是否存在尚未打印的新解。"""
        return bool(self._pending_solutions)

    # -----------------------------
    # print（仅输出，不新增记录）
    # -----------------------------
    def print_pending(
        self,
        question: Question,
        question_rule: _QuestionRule,
    ) -> bool:
        """打印自上次调用以来新加入的解（含等价扩展），并根据配置决定是否终止推理。

        :return: 是否应该终止推理
        """
        if not self._pending_solutions:
            return False

        question_str = ", ".join(str(q) for q in question.question)

        while self._pending_solutions:
            # HACK: 使用 pop() 逆序输出新增解；当前没有需求时不特地排序或引入 deque 优化。
            combination = self._pending_solutions.pop()
            key = self._solution_key(combination)

            # pending 理论上都未打印过，但这里仍做一次防御性处理
            if key in self._printed_solution_keys:
                continue

            self._printed_solution_keys.add(key)

            if combination:
                # 输出时除了 matcher 找到的绑定，也包含等价扩展后的绑定。
                binding_str = self._display_binding(combination)
                logger.result("Query solution: %s variable bindings: %s", question_str, binding_str)  # type: ignore[attr-defined]
            else:
                logger.result("Query solution: %s", question_str)  # type: ignore[attr-defined]
                if not question_rule.free_variables:
                    return True

            if self._should_terminate(self.args.interactive_query_mode):
                return True

        self._pending_solutions.clear()
        return False

    # -----------------------------
    # backward compatible wrapper
    # -----------------------------
    def print_and_collect(
        self,
        question: Question,
        solutions: list[Mapping[Variable, Constant | CompoundTerm]],
        question_rule: _QuestionRule,
        *,
        equivalence: Equivalence | None = None,
    ) -> bool:
        """兼容旧接口：先 add 再 print。"""
        self.add_solutions(solutions, equivalence=equivalence)
        return self.print_pending(question, question_rule)

    def _record_with_equivalents(
        self,
        solution: Mapping[Variable, Constant | CompoundTerm],
        *,
        equivalence: Equivalence | None,
    ) -> None:
        """记录 solution 及其等价扩展解（去重），不做打印。"""
        self._register_solution(solution)

        # interactive_query_mode=first 时仅输出第一个解，不做等价扩展（add_solutions 同样走该逻辑）。
        if self.args.interactive_query_mode == "first":
            self._solution_count = len(self._seen_solution_keys)
            return

        # 对每个绑定值做等价扩展：逐变量替换。
        if equivalence is not None and solution:
            items = sorted(solution.items(), key=lambda kv: kv[0].display_name)
            for var, term in items:
                for equiv_term in equivalence.get_equiv_item(term):
                    expanded = dict(solution)
                    expanded[var] = equiv_term  # 这里开销会比较大，不过早晚得承担。需要优化的话可以挪到打印solution时候再干
                    self._register_solution(expanded)

        self._solution_count = len(self._seen_solution_keys)

    def _register_solution(self, solution: Mapping[Variable, Constant | CompoundTerm]) -> bool:
        """将解注册到 seen/pending/collected（去重后）。"""
        key = self._solution_key(solution)
        if key in self._seen_solution_keys:
            return False

        self._seen_solution_keys.add(key)

        # 只把“从未打印过”的解放入 pending
        if key not in self._printed_solution_keys:
            self._pending_solutions.append(dict(solution))

        if self.args.save_solutions:
            self.collected_solutions.append(dict(solution))

        return True

    @staticmethod
    def _solution_key(solution: Mapping[Variable, Constant | CompoundTerm]) -> tuple[tuple[str, str], ...]:
        """用于去重的稳定 key：按变量 display_name 排序后，取 (var_display_name, term_str)。"""
        if not solution:
            return ()
        items = sorted(solution.items(), key=lambda kv: kv[0].display_name)
        return tuple((var.display_name, str(term)) for var, term in items)

    @staticmethod
    def _display_binding(combination: Mapping[Variable, Constant | CompoundTerm]) -> str:
        entries: list[str] = []
        for var, term in sorted(combination.items(), key=lambda kv: kv[0].display_name):
            entries.append(f"{var.display_name}={term}")
        return ", ".join(entries)

    @staticmethod
    def _should_terminate(mode: Literal["interactive", "first", "all"]) -> bool:
        if mode == "first":
            logger.info("Configured to output only the first query solution; stopping inference.")
            return True
        if mode == "all":
            return False

        user_input = input("发现解，输入 ';' 并回车继续本次推理；输入其他任意键并回车将终止本次推理: ").strip()  # TODO: input不太适合后续扩展
        if user_input != ';':
            logger.info("用户选择终止推理。")
            return True
        return False

    def get_all_solutions(self) -> list[Mapping[Variable, Constant | CompoundTerm]]:
        """获取所有的solutions"""
        return self.collected_solutions

    def get_solution_count(self) -> int:
        """获取solution的数目"""
        return self._solution_count


class StatusChecker:
    """通用状态检查逻辑"""

    def __init__(self, equivalence: Equivalence, sk_system_handler: SankuManagementSystem):
        self.equivalence = equivalence
        self.sk_system_handler = sk_system_handler

    def check_conflict(self, new_facts: list[FACT_TYPE]) -> bool:
        """检查是否发生了矛盾。!!! 注意这里传入的是new_facts"""
        return self._has_conflict_occurred(new_facts)

    @staticmethod
    def _has_conflict_occurred(new_facts: list[FACT_TYPE]) -> bool:
        """检查是否发生了矛盾"""
        # TODO: 实现矛盾检测逻辑
        return False


class MainLoopManager:
    """主循环管理器"""

    def __init__(self, status_checker: StatusChecker, args: RunControlConfig | None = None):
        self.status_checker = status_checker
        self.args = args or RunControlConfig()
        self._current_iteration = 0

        self.normal_rule_activated: dict[Rule, bool] = {}  # 如果某条rule上一轮有新事实生成，则True；反之False
        self._true_count = -1

    def check_status(self,
                     current_facts: list[FACT_TYPE],
                     question: Question) -> InferenceStatus:
        """检查主循环状态"""
        # 先检查迭代次数
        if self._current_iteration >= self.args.iteration_limit:
            return InferenceStatus.MAX_ITERATIONS_REACHED

        if self._true_count == 0:
            return InferenceStatus.FIXPOINT_REACHED

        if self._current_iteration == 0 and question.question and all(q in current_facts for q in question.question):
            # FIXME: 这里的判断可能需要分的更细致
            return InferenceStatus.SUCCESS

        return InferenceStatus.CONTINUE

    def next_iteration(self) -> None:
        """进入下一轮迭代"""
        self._current_iteration += 1

    def reset(self) -> None:
        """重置计数"""
        self._current_iteration = 0

    def initial_manager(self, normal_rules: Sequence[Rule] | None = None, *, resume: bool = False) -> None:
        """
        修改current rules, 为当前一个question的推理做准备
        :raise ValueError: 如果待推理问题不变，仅中止引擎并重新推理时，认为不需要修改rules（其他各处也应调整）
        """
        if normal_rules is None:
            if resume:
                self.normal_rule_activated = dict.fromkeys(self.normal_rule_activated, True)  # 继续推理当前时，使用之前存储好的rule
                # 但是清空之前的_activated记录（因为可能有新事实变化）
            else:
                raise ValueError("normal_rules is None")
        else:
            self.normal_rule_activated = dict.fromkeys(normal_rules, True)  # 如果某条rule上一轮有新事实生成，则True；反之False
        self._true_count = len(self.normal_rule_activated)

    @property
    def iteration(self) -> int:
        """获取当前迭代次数"""
        return self._current_iteration

    def update_normal_rule_activation(self, new_facts: list[FACT_TYPE], used_rule: Rule) -> None:
        """每条规则完成推理后，使用本函数更新main_manager，用于判断是否所有的规则同时达到了不动点"""
        activated = bool(new_facts)

        old_value = self.normal_rule_activated.get(used_rule)
        self.normal_rule_activated[used_rule] = activated

        if old_value and not activated:
            self._true_count -= 1
        elif not old_value and activated:
            self._true_count += 1

    def is_at_fixpoint(self) -> bool:
        """是否所有 normal rule 都已不再产生新事实（不动点，仅看 normal rules）"""
        return self._true_count == 0


class ExecutorManager:
    """执行器管理器"""

    def __init__(self, status_checker: StatusChecker, solution_manager: QuerySolutionManager,
                 args: ExecutorConfig):
        self.status_checker = status_checker
        self.args = args
        self._current_step = 1
        self.solution_manager = solution_manager

    def check_status(self, new_facts: list[FACT_TYPE], question: Question,
                     solutions: list[Mapping[Variable, Constant | CompoundTerm]],
                     question_rule: _QuestionRule | None) -> InferenceStatus:
        """检查执行器状态"""
        if self.args.executing_max_steps != -1 and self._current_step >= self.args.executing_max_steps:
            return InferenceStatus.MAX_STEPS_REACHED  # 有解是success，且需要解

        # 检查是否有矛盾
        if self.status_checker.check_conflict(new_facts):  # TODO: 冲突时同步清理 solutions
            return InferenceStatus.CONFLICT_DETECTED

        if question_rule is not None:
            # 1) 记录 matcher 找到的解，并同步记录其等价元素
            if solutions:
                completed_solutions = self._complete_question_solutions(question, solutions)
                self.solution_manager.add_solutions(
                    completed_solutions,
                    equivalence=self.status_checker.equivalence,
                )

            # 2) 仅打印尚未输出过的新解
            if self.solution_manager.has_pending_solutions():
                terminate_sign = self.solution_manager.print_pending(question, question_rule)
                if terminate_sign:
                    return InferenceStatus.SUCCESS

        return InferenceStatus.CONTINUE

    @staticmethod
    def _complete_question_solutions(
        question: Question,
        solutions: list[Mapping[Variable, Constant | CompoundTerm]],
    ) -> list[Mapping[Variable, Constant | CompoundTerm]]:
        if not question.free_variables:
            return solutions

        filler = _QuestionRule.QUESTION_VAR_FILLER
        filled_solutions: list[Mapping[Variable, Constant | CompoundTerm]] = []

        for solution in solutions:
            known_vars = [k.display_name for k in solution]
            missing_vars = [var for var in question.free_variables if var.display_name not in known_vars]
            # HACK: 由于question变更为question rule时候进行了重命名（且由于可能被DNF split的缘故，这个重命名似乎是必要的）
            # 系统中question和question rule的Variable instance之间产生了不一致。这里做一层初步修复
            # 这里会遗留一个问题是，solution里面的Variable instance其实是不一样的，只是打印的display name一样
            filled_solutions.append({**solution, **dict.fromkeys(missing_vars, filler)})

        return filled_solutions

    def next_step(self) -> None:
        """执行下一步"""
        self._current_step += 1

    def reset_for_new_inference(self) -> None:
        """为新的推理过程重置步数计数（仅在开始全新推理时调用）"""
        self._current_step = 1

    @property
    def step_num(self) -> int:
        """获取当前步数"""
        return self._current_step


def create_main_loop_manager(equivalence: Equivalence,
                             sk_system_handler: SankuManagementSystem,
                             args: RunControlConfig | None = None) -> MainLoopManager:
    """创建主循环管理器"""
    status_checker = StatusChecker(equivalence, sk_system_handler)
    return MainLoopManager(status_checker, args)


def create_executor_manager(equivalence: Equivalence,
                            sk_system_handler: SankuManagementSystem,
                            solution_manager: QuerySolutionManager,
                            args: ExecutorConfig) -> ExecutorManager:
    """创建执行器管理器"""
    status_checker = StatusChecker(equivalence, sk_system_handler)
    return ExecutorManager(status_checker, solution_manager, args)
