from pyparsing import ParseResults
from rapt2.treebrd.condition_node import (
    IdentityConditionNode,
    BinaryConditionNode,
    BinaryConditionalOperator,
    ConditionNode,
    UnaryConditionNode,
    UnaryConditionalOperator,
)
from .schema import Schema
from .node import (
    SelectNode,
    ProjectNode,
    RenameNode,
    AssignNode,
    CrossJoinNode,
    NaturalJoinNode,
    UnionNode,
    DifferenceNode,
    IntersectNode,
    ThetaJoinNode,
    FullOuterJoinNode,
    LeftOuterJoinNode,
    RightOuterJoinNode,
    RelationNode,
)
from .grammars.proto_grammar import ProtoGrammar


class TreeBRD:
    """
    A Tree Builder for Relational (Algebra) Data. TreeBRD is a factory class
    that builds forests of relational algebra syntax trees. STARBuilder
    """

    grammar: ProtoGrammar

    def __init__(self, grammar: ProtoGrammar):
        self.grammar = grammar

    def build(self, instring, schema):
        ra = self.grammar.parse(instring)
        _schema = Schema(schema)
        return [self.to_node(statement, _schema) for statement in ra[:]]

    def to_node(self, exp: ParseResults, schema):
        """
        Return a Node that is the root of the parse tree for the the specified
        expression.

        :param exp: A list that represents a relational algebra expression.
        Assumes that this list was generated by pyparsing.
        :param schema: A dictionary of relation names to attribute names used
        for verification and generating attributes.
        :return: A Node.
        """

        # A relation node.
        if len(exp) == 1 and isinstance(exp[0], str):
            name = exp[0]
            node = RelationNode(name=name, schema=schema)

        # An expression.
        elif len(exp) == 1 and isinstance(exp[0], ParseResults):
            node = self.to_node(exp[0], schema)

        # Unary operators.
        elif isinstance(exp[0], str) and self.grammar.is_unary(exp[0]):
            child = self.to_node(exp[2:], schema)
            node = self.create_unary_node(
                operator=exp[0], child=child, param=exp[1], schema=schema
            )

        # Assignment.
        elif exp[1] is self.grammar.syntax.assign_op:
            child = self.to_node(exp[2:], schema)
            node = self.create_unary_node(
                operator=exp[1], child=child, param=exp[0], schema=schema
            )

        # Binary operators.
        elif self.grammar.is_binary(exp[1]):
            # Pyparsing will put different operators with the same precedence
            # in the same list. This can be a problem when we mix operators with
            # and without parameters (for example join). We avoid this below
            # and build from right to left, to create the correct syntax tree.
            if isinstance(exp[-2], str):
                # Operator without parameters
                op_pos = -2
                param = None
            else:
                op_pos = -3
                param = exp[-2]

            operator = exp[op_pos]
            left = self.to_node(exp[:op_pos], schema)
            right = self.to_node(exp[-1], schema)
            node = self.create_binary_node(
                operator=operator, left=left, right=right, param=param
            )

        else:
            raise ValueError

        return node

    def create_condition_node(self, conditions: ParseResults) -> ConditionNode:
        if isinstance(conditions, str):
            # Attribute Condition Node
            return IdentityConditionNode(conditions)

        elif len(conditions) == 2:
            # Unary Condition Node
            op, right = conditions
            right_node = self.create_condition_node(right)

            return UnaryConditionNode(
                op=UnaryConditionalOperator.from_syntax(self.grammar.syntax, op),
                child=right_node,
            )

        elif len(conditions) == 3:
            # Binary Condition Node
            left, op, right = conditions
            left_node = self.create_condition_node(left)
            right_node = self.create_condition_node(right)

            return BinaryConditionNode(
                op=BinaryConditionalOperator.from_syntax(self.grammar.syntax, op),
                left=left_node,
                right=right_node,
            )

        raise ValueError

    def create_unary_node(
        self, operator, child, param: ParseResults | None = None, schema=None
    ):
        """
        Return a Unary Node whose type depends on the specified operator.

        :param schema:
        :param child:
        :param operator: A relational algebra operator (see constants.py)
        :param param: A list of parameters for the operator.
        :return: A Unary Node.
        """

        if operator == self.grammar.syntax.select_op:
            condition = self.create_condition_node(param[0])
            node = SelectNode(child, condition)

        elif operator == self.grammar.syntax.project_op:
            node = ProjectNode(child, param)

        elif operator == self.grammar.syntax.rename_op:
            name = None
            attributes = []
            if isinstance(param[0], str):
                name = param.pop(0)
            if param:
                attributes = param[0]
            node = RenameNode(child, name, attributes, schema)

        elif operator == self.grammar.syntax.assign_op:
            name = param[0]
            attributes = [] if len(param) < 2 else param[1]
            node = AssignNode(child, name, attributes, schema)

        else:
            raise ValueError

        return node

    def create_binary_node(
        self, operator, left, right, param: ParseResults | None = None
    ):
        """
        Return a Node whose type depends on the specified operator.

        :param operator: A relational algebra operator (see constants.py)
        :return: A Node.
        """

        # Join operators
        if operator == self.grammar.syntax.join_op:
            node = CrossJoinNode(left, right)

        elif operator == self.grammar.syntax.natural_join_op:
            node = NaturalJoinNode(left, right)

        elif operator == self.grammar.syntax.theta_join_op:
            condition = self.create_condition_node(param[0])
            node = ThetaJoinNode(left, right, condition)

        elif operator == self.grammar.syntax.full_outer_join_op:
            condition = self.create_condition_node(param[0])
            node = FullOuterJoinNode(left, right, condition)

        elif operator == self.grammar.syntax.left_outer_join_op:
            condition = self.create_condition_node(param[0])
            node = LeftOuterJoinNode(left, right, condition)

        elif operator == self.grammar.syntax.right_outer_join_op:
            condition = self.create_condition_node(param[0])
            node = RightOuterJoinNode(left, right, condition)

        # Set operators
        elif operator == self.grammar.syntax.union_op:
            node = UnionNode(left, right)

        elif operator == self.grammar.syntax.difference_op:
            node = DifferenceNode(left, right)

        elif operator == self.grammar.syntax.intersect_op:
            node = IntersectNode(left, right)

        else:
            raise ValueError

        return node
