from __future__ import annotations

from typing import TYPE_CHECKING, overload

from kele.syntax import Constant, Variable
from kele.syntax.base_classes import Assertion, CompoundTerm, Formula, Rule
from kele.syntax.dnf_converter import RuleSafetyProcesser

if TYPE_CHECKING:
    from collections.abc import Sequence


class RulePreprocessor:
    """Shared preprocessing pipeline for rule normalization.

    This encapsulates splitting and variable-collision resolution so callers can
    reuse the same behavior without creating another RuleBase.
    """

    def __init__(self, *, start_rule_suffix: int = 1, start_var_counter: int = 0) -> None:
        self.rule_suffix = start_rule_suffix
        self._var_counter = start_var_counter

    def preprocess_rule[T1: Rule](self, rule: T1) -> list[T1]:
        """Split, rename, and collision-resolve a single rule."""
        return [self._resolve_variable_collisions(single_rule) for single_rule in self._split_and_rename_rule(rule)]

    def preprocess_rules[T1: Rule](self, rules: Sequence[T1]) -> list[T1]:
        """Apply preprocessing to each rule and flatten the outputs."""
        processed: list[T1] = []
        for rule in rules:
            processed.extend(self.preprocess_rule(rule))
        return processed

    def _resolve_variable_collisions[T1: Rule](self, rule: T1) -> T1:
        """Rename variables in a rule to avoid collisions across rules."""
        var_map: dict[str, Variable] = {}
        new_head = self._generate_renamed_item(rule.head, var_map)
        new_body = self._generate_renamed_item(rule.body, var_map)

        return rule.replace(head=new_head, body=new_body)

    def _split_and_rename_rule[T1: Rule](self, rule: T1) -> list[T1]:
        """Assign names to unnamed rules and split rules into DNF-safe units."""
        origin_name = rule.name or f"rule_{self.rule_suffix}"
        self.rule_suffix += 1

        rule_spliter = RuleSafetyProcesser()
        splited_rules = rule_spliter.split_rule_and_process_safety(rule)

        if len(splited_rules) == 1:
            return [splited_rules[0].replace(name=origin_name)]

        result: list[T1] = []
        for idx, single_rule in enumerate(splited_rules):
            new_name = f"{origin_name}_{idx + 1}"
            result.append(single_rule.replace(name=new_name))

        return result

    @overload
    def _generate_renamed_item(self, item: Variable, var_map: dict[str, Variable]) -> Variable: ...
    @overload
    def _generate_renamed_item(self, item: Constant, var_map: dict[str, Variable]) -> Constant: ...
    @overload
    def _generate_renamed_item(self, item: CompoundTerm, var_map: dict[str, Variable]) -> CompoundTerm: ...
    @overload
    def _generate_renamed_item(self, item: Assertion, var_map: dict[str, Variable]) -> Assertion: ...
    @overload
    def _generate_renamed_item(self, item: Formula, var_map: dict[str, Variable]) -> Formula: ...
    @overload
    def _generate_renamed_item(self, item: None, var_map: dict[str, Variable]) -> None: ...

    def _generate_renamed_item(
        self,
        item: Formula | Assertion | CompoundTerm | Variable | Constant | None,
        var_map: dict[str, Variable],
    ) -> Formula | Assertion | CompoundTerm | Variable | Constant | None:
        if isinstance(item, Variable):
            if item.symbol not in var_map:
                new_name = f"_v{self._var_counter}"
                self._var_counter += 1
                var_map[item.symbol] = item.create_renamed_variable(new_name)
                # risk: 两条规则里如果都用到了x，然后使用者使用了同一个x的instance，
                # 我们是否有理由认为它俩是一个x？（我认为不合理，所以我们应当不管instance，而只是取它的value本身。但这样会不会给使用者带来困惑？
                # 他可能刻意用了同一个instance，甚至在里面存了一些自定义的信息。但我们可能是直接丢弃然后换个了新名字新地址）
            return var_map[item.symbol]

        if isinstance(item, Constant):
            return item

        if isinstance(item, CompoundTerm):
            return CompoundTerm.from_parts(
                item.operator,
                tuple(self._generate_renamed_item(arg, var_map) for arg in item.arguments),
                description=item.description,
            )

        if isinstance(item, Assertion):
            return Assertion.from_parts(
                self._generate_renamed_item(item.lhs, var_map),
                self._generate_renamed_item(item.rhs, var_map),
                description=item.description,
            )

        if isinstance(item, Formula):
            return Formula.from_parts(
                self._generate_renamed_item(item.formula_left, var_map),
                item.connective,
                self._generate_renamed_item(item.formula_right, var_map),
                description=item.description,
            )

        if item is None:
            return None

        raise ValueError(f"Unknown item type: {type(item)}")
