import re
from typing import List, Tuple

from statemachine import State, StateMachine
from statemachine.exceptions import TransitionNotAllowed


class FSM(StateMachine):
    def __init__(self):
        super().__init__()
        self._stack = []

    def process(self, node: dict):
        phase = node["_phase"]
        if phase == "enter":
            self.enter(node)

        name = re.sub(r"(?<!^)(?=[A-Z])", "_", node["_type"]).lower()
        if name in ["select_stmt", "update_stmt", "insert_stmt", "delete_stmt"]:
            name = "stmt"
        event = phase + "_" + name
        if event in self._events:
            try:
                self.send(event, node=node)
            except TransitionNotAllowed:
                pass

        if phase == "leave":
            self.leave(node)

    def enter(self, node: dict):
        self._stack.append(node)

    def leave(self, node: dict):
        n = self._stack.pop()
        tp = n["_type"]
        assert (
            tp == node["_type"]
        ), f"enter type & leave type mismatch: {n['_type']} vs {node['_type']}"
        return n


class CorrelationFSM(FSM):
    init = State(initial=True)
    stmt_end = State("End", final=True)
    on_condition = State("On Condition")
    binary_operation = State("Binary Operation")
    receive_column = State("Receiving Column")

    leave_stmt = init.to(stmt_end, cond="outer_stmt")

    enter_on_condition = init.to(on_condition)
    leave_on_condition = on_condition.to(init)

    enter_binary_operation_expr = on_condition.to(binary_operation, cond="op_eq")
    leave_binary_operation_expr = binary_operation.to(on_condition, cond="op_eq")

    enter_column_name = binary_operation.to(receive_column)
    leave_column_name = receive_column.to(binary_operation)

    def __init__(self):
        super().__init__()
        self._correlations = []
        self._columns = []

    def op_eq(self, node: dict):
        if "op" in node:
            return node["op"] == "eq"
        else:
            return self._stack[-1]["op"] == "eq"

    def outer_stmt(self, node: dict):
        return len(self._stack) == 1

    @receive_column.enter
    def save_column_name(self, node: dict):
        if node["table"]:
            self._columns.append(node["table"] + "." + node["name"])
        else:
            self._columns.append(node["name"])

    @binary_operation.exit
    def save_correlation(self, node: dict, target: State):
        if target.id == "on_condition":
            self._correlations.append((self._columns[0], self._columns[1]))
            self._columns = []

    def correlations(self) -> List[Tuple[str, str]]:
        return self._correlations


class MarkerFSM(FSM):
    init = State(initial=True)
    stmt_end = State("End", final=True)
    receive_marker_column = State("Receiving Marker Column")
    receive_assignment_column = State("Receiving Assignment Column")
    discard_column = State("Discarding Column")
    binary_operation = State("Binary Operation")
    assignment = State("Assignment")
    count_assignment_marker = State("Count Assignment Marker")
    count_marker = State("Count Marker")

    leave_stmt = init.to(stmt_end, cond="outer_stmt")

    enter_binary_operation_expr = init.to(binary_operation, unless="logic_op")
    leave_binary_operation_expr = count_marker.to(init, unless="logic_op")

    enter_param_marker_expr = (
        count_marker.to.itself() | count_assignment_marker.to.itself()
    )

    enter_assignment = init.to(assignment)
    leave_assignment = count_assignment_marker.to(init)

    enter_column_name = (
        init.to(discard_column)
        | binary_operation.to(receive_marker_column)
        | assignment.to(receive_assignment_column)
    )
    leave_column_name = (
        discard_column.to(init)
        | receive_marker_column.to(count_marker)
        | receive_assignment_column.to(count_assignment_marker)
    )

    def __init__(self):
        super().__init__()
        self._marker_columns = []
        self._assignment_columns = []
        self._marker_counts = []

    def assignments(self) -> List[str]:
        assert self.current_state.final
        column_names = []
        for i in range(len(self._assignment_columns)):
            count = self._marker_counts[i]
            for _ in range(count):
                column_names.append(self._assignment_columns[i])
        return column_names

    def markers(self) -> List[str]:
        assert self.current_state.final
        column_names = []
        for i in range(len(self._marker_columns)):
            count = self._marker_counts[i + len(self._assignment_columns)]
            for _ in range(count):
                column_names.append(self._marker_columns[i])
        return column_names

    @receive_marker_column.enter
    def save_marker_column_name(self, node: dict):
        if node["table"]:
            self._marker_columns.append(node["table"] + "." + node["name"])
        else:
            self._marker_columns.append(node["name"])
        self._marker_counts.append(0)

    @receive_assignment_column.enter
    def save_assignment_column_name(self, node: dict):
        if node["table"]:
            self._assignment_columns.append(node["table"] + "." + node["name"])
        else:
            self._assignment_columns.append(node["name"])
        self._marker_counts.append(0)

    @enter_param_marker_expr.on
    def inc_count_marker(self):
        self._marker_counts[-1] += 1

    def logic_op(self, node: dict):
        if "op" in node:
            return node["op"] in ["and", "or"]
        else:
            return self._stack[-1]["op"] in ["and", "or"]

    def outer_stmt(self, node: dict):
        return len(self._stack) == 1