from typing import Any, NamedTuple

from pyparsing import ParseResults

from .condition_node import (
    BinaryConditionalOperator,
    BinaryConditionNode,
    ConditionNode,
    IdentityConditionNode,
    UnaryConditionalOperator,
    UnaryConditionNode,
)

from .grammars.proto_grammar import ProtoGrammar
from .node import (
    Node,
    AssignNode,
    CrossJoinNode,
    DifferenceNode,
    FullOuterJoinNode,
    FunctionalDependencyNode,
    InclusionEquivalenceNode,
    InclusionSubsumptionNode,
    IntersectNode,
    LeftOuterJoinNode,
    MultivaluedDependencyNode,
    NaturalJoinNode,
    PrimaryKeyNode,
    ProjectNode,
    RelationNode,
    RenameNode,
    RightOuterJoinNode,
    SelectNode,
    ThetaJoinNode,
    UnionNode,
)
from .schema import Schema


class _BinaryParts(NamedTuple):
    """Components of a flattened binary expression extracted from ParseResults."""

    left: ParseResults
    operator: str
    operator_params: Any
    right: Any


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):
        """
        :param grammar: the grammar used to parse relational algebra strings.
        """
        self.grammar = grammar
        self._conditional_join_types = {
            grammar.syntax.theta_join_op: ThetaJoinNode,
            grammar.syntax.full_outer_join_op: FullOuterJoinNode,
            grammar.syntax.left_outer_join_op: LeftOuterJoinNode,
            grammar.syntax.right_outer_join_op: RightOuterJoinNode,
        }
        self._unary_dependency_types = {
            grammar.syntax.mvd_op: MultivaluedDependencyNode,
            grammar.syntax.fd_op: FunctionalDependencyNode,
        }

    def build(self, instring, schema) -> list[Node]:
        """
        Parse a relational algebra string and build a forest of syntax trees.

        :param instring: a relational algebra string (semicolon-separated
            statements).
        :param schema: a dict mapping relation names to attribute name lists.
        :return: a list of root nodes, one per statement.

        .. note::
            Assignment statements (e.g. ``R := …``) mutate *schema* as a
            side effect, registering the new relation so that later statements
            can reference it.
        """
        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.
        """

        # Dependency statements (pk, mvd, fd, inc) start with a dependency
        # operator keyword — check this first to avoid falling into RA dispatch.
        if self.is_dependency_statement(exp):
            return self.create_dependency_node(exp, schema)

        # Single string → bare relation name from schema.
        if len(exp) == 1 and isinstance(exp[0], str):
            name = exp[0]
            node = RelationNode(name=name, schema=schema)

        # Single nested ParseResults → parenthesised sub-expression.
        elif len(exp) == 1 and isinstance(exp[0], ParseResults):
            node = self.to_node(exp[0], schema)

        # Leading string that is a unary operator → [op, params, child…].
        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, operator_params=exp[1], schema=schema
            )

        # Second element is the assignment operator → [lhs, :=, child…].
        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, operator_params=exp[0], schema=schema
            )

        # Second element is a binary operator → [left, op, (param?,) right].
        elif self.grammar.is_binary(exp[1]):
            parts = self._extract_binary_parts(exp)
            left = self.to_node(parts.left, schema)
            right = self.to_node(parts.right, schema)
            node = self.create_binary_node(
                operator=parts.operator,
                left=left,
                right=right,
                operator_params=parts.operator_params,
            )

        else:
            raise ValueError(f"Unrecognised expression: {list(exp)}")

        return node

    @staticmethod
    def _extract_binary_parts(exp: ParseResults) -> _BinaryParts:
        """Split a flattened binary expression into (left, operator, param, right).

        Pyparsing flattens same-precedence operators into one list, e.g.:
            "A \\join B \\union C" → [A, \\join, B, \\union, C]
        We process from right to left so the rightmost operator binds first.
        The rightmost operand is always exp[-1]. The operator position
        depends on whether it has parameters:
            Without params: [..., operator, right]       → op at -2
            With params:    [..., operator, param, right] → op at -3
        """
        if isinstance(exp[-2], str):
            op_pos = -2
            param = None
        else:
            op_pos = -3
            param = exp[-2]

        return _BinaryParts(
            left=exp[:op_pos],
            operator=exp[op_pos],
            operator_params=param,
            right=exp[-1],
        )

    def create_condition_node(self, conditions: ParseResults) -> ConditionNode:
        """
        Recursively build a ConditionNode tree from parsed condition tokens.

        :param conditions: pyparsing result representing a condition expression.
        :return: the root ConditionNode of the condition tree.
        :raises ValueError: if the conditions structure is unrecognised.
        """
        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(
            f"Unrecognised condition structure (length {len(conditions)}): "
            f"{list(conditions)}"
        )

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

        :param operator: A relational algebra operator (see constants.py)
        :param child: The child node.
        :param operator_params: Parsed parameters whose shape depends on the
            operator — conditions for select, attribute lists for project,
            name/attributes for rename and assign.
        :param schema: Schema for relation validation.
        :return: A Unary Node.
        """

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

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

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

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

        else:
            raise ValueError(f"Unknown unary operator: {operator}")

        return node

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

        :param operator: A relational algebra operator (see constants.py)
        :param left: Left child node.
        :param right: Right child node.
        :param operator_params: Parsed parameters (e.g. join conditions).
        :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 in self._conditional_join_types:
            condition = self.create_condition_node(operator_params[0])
            node = self._conditional_join_types[operator](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(f"Unknown binary operator: {operator}")

        return node

    def is_dependency_statement(self, exp: ParseResults) -> bool:
        """
        Check if the expression is a dependency statement.

        :param exp: A parsed expression
        :return: True if it's a dependency statement
        """
        if not exp or len(exp) < 2:
            return False

        return exp[0] in self.grammar.syntax.dependency_operators

    def create_dependency_node(self, exp: ParseResults, schema):
        """
        Create a dependency node from a parsed dependency statement.

        :param exp: A parsed dependency expression
        :param schema: A schema object
        :return: A dependency node
        """
        operator = exp[0]

        if operator == self.grammar.syntax.pk_op:
            # pk_{attributes} relation
            attributes = exp[1]
            relation_name = exp[2]
            return PrimaryKeyNode(relation_name, attributes)

        elif operator in (self.grammar.syntax.mvd_op, self.grammar.syntax.fd_op):
            # mvd_{attribute1, attribute2} relation OR mvd_{attribute1, attribute2} \select_{conditions} relation
            # fd_{attribute1, attribute2} relation OR fd_{attribute1, attribute2} \select_{conditions} relation
            return self._create_dependency_with_optional_select(exp, schema, operator)

        elif operator == self.grammar.syntax.inc_equiv_op:
            return self._create_inclusion_node(exp, schema, InclusionEquivalenceNode)

        elif operator == self.grammar.syntax.inc_subset_op:
            return self._create_inclusion_node(exp, schema, InclusionSubsumptionNode)

        else:
            raise ValueError(f"Unknown dependency operator: {operator}")

    def _create_dependency_with_optional_select(self, exp, schema, operator):
        """
        Create a dependency node (MVD or FD) with optional select conditions.

        Parse layout:
            Simple:      [op, attributes, relation]
            With select: [op, attributes, \\select, conditions, relation]

        :param exp: Parsed expression containing attributes and optional select
        :param schema: Schema for relation validation
        :param operator: The dependency operator (mvd_op or fd_op)
        :return: MultivaluedDependencyNode or FunctionalDependencyNode
        """
        node_class = self._unary_dependency_types[operator]
        attributes = exp[1]

        if len(exp) == 3:
            # Simple: op_{attrs} relation
            relation_name = exp[2]
            relation_node = RelationNode(relation_name, schema)
        else:
            # With select: op_{attrs} \select_{cond} relation
            _, conditions, relation_name = exp[2], exp[3], exp[4]
            condition_node = self.create_condition_node(conditions[0])
            base_relation = RelationNode(relation_name, schema)
            relation_node = SelectNode(base_relation, condition_node)

        return node_class(relation_name, attributes, relation_node)

    def _create_inclusion_node(self, exp, schema, node_class):
        """
        Create an inclusion dependency node (equivalence or subsumption).

        :param exp: Parsed expression containing attributes and relation expressions
        :param schema: Schema for relation validation
        :param node_class: InclusionEquivalenceNode or InclusionSubsumptionNode
        :return: An inclusion dependency node
        """
        attributes = exp[1]
        relations_expr = exp[2]

        if isinstance(relations_expr[0], str):
            # Simple form: inc_{attributes} (relation1, relation2)
            relation_names = relations_expr
            left_child = RelationNode(relation_names[0], schema)
            right_child = RelationNode(relation_names[1], schema)
        else:
            # Complex form with select: inc_{attributes} (select_or_relation, select_or_relation)
            left_expr = relations_expr[0]
            right_expr = relations_expr[1]
            left_child = self._create_cond_dep_expr_node(left_expr, schema)
            right_child = self._create_cond_dep_expr_node(right_expr, schema)
            relation_names = [left_child.name, right_child.name]
        return node_class(relation_names, attributes, left_child, right_child)

    def _create_cond_dep_expr_node(self, expr, schema):
        """
        Create a node from a cond_dep_expr (either a bare relation or a select over a relation).

        Parse layout:
            Bare relation:  "relation_name" or ["relation_name"]
            With select:    [\\select, conditions, relation_name]

        :param expr: Parsed expression
        :param schema: Schema for relation validation
        :return: RelationNode or SelectNode
        """
        if isinstance(expr, str):
            return RelationNode(expr, schema)
        elif len(expr) == 1 and isinstance(expr[0], str):
            return RelationNode(expr[0], schema)
        else:
            _, conditions, relation_name = expr[0], expr[1], expr[2]
            condition_node = self.create_condition_node(conditions[0])
            base_relation = RelationNode(relation_name, schema)
            return SelectNode(base_relation, condition_node)
