"""
Copyright 2024, Adnan Khan

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import re


class Node:
    def __init__(self, type, value=None, children=None):
        self.type = type
        self.value = value
        self.children = children or []


class ExpressionParser:
    def __init__(self, expression):
        self.expression = expression
        self.tokens = self.tokenize(expression)
        self.parsed_expression = self.parse_expression(self.tokens)

    def __print_ast(self, node, level=0):
        indent = "  " * level
        print(f"{indent}{node.type}: {node.value}")
        for child in node.children:
            self.__print_ast(child, level + 1)

    @staticmethod
    def tokenize(expression: str):
        token_specification = [
            ("LPAREN", r"\("),
            ("RPAREN", r"\)"),
            ("COMMA", r","),
            ("AND", r"&&"),
            ("OR", r"\|\|"),
            ("EQUALS", r"=="),
            ("NOT_EQUALS", r"!="),
            ("NEGATION", r"!"),  # Unary negation
            ("IDENTIFIER", r"[a-zA-Z_.][a-zA-Z0-9_.*]*"),
            ("STRING", r"'[^']*'"),
            ("WHITESPACE", r"\s+"),
        ]
        tok_regex = "|".join(
            f"(?P<{name}>{pattern})" for name, pattern in token_specification
        )
        return [
            (match.lastgroup, match.group(0))
            for match in re.finditer(tok_regex, expression)
            if match.lastgroup != "WHITESPACE"
        ]

    @staticmethod
    def expect_token(tokens, expected_type):
        if tokens and tokens[0][0] == expected_type:
            return tokens.pop(0)
        raise SyntaxError(
            f"Expected {expected_type} but got {tokens[0] if tokens else 'no more tokens'}"
        )

    @staticmethod
    def parse_expression(tokens):
        if not tokens:
            raise SyntaxError("Empty expression")
        node = ExpressionParser.parse_logical_or(tokens)
        return node

    @staticmethod
    def parse_logical_or(tokens):
        node = ExpressionParser.parse_logical_and(tokens)
        while tokens and tokens[0][0] == "OR":
            op = tokens.pop(0)
            rhs = ExpressionParser.parse_logical_and(tokens)
            node = Node("logical_or", op[1], [node, rhs])
        return node

    @staticmethod
    def parse_logical_and(tokens):
        node = ExpressionParser.parse_comparison(tokens)
        while tokens and tokens[0][0] == "AND":
            op = tokens.pop(0)
            rhs = ExpressionParser.parse_comparison(tokens)
            node = Node("logical_and", op[1], [node, rhs])
        return node

    @staticmethod
    def parse_comparison(tokens):
        node = ExpressionParser.parse_unary(
            tokens
        )  # Updated from parse_primary to parse_unary
        while tokens and tokens[0][0] in ("EQUALS", "NOT_EQUALS"):
            op = tokens.pop(0)
            rhs = ExpressionParser.parse_unary(
                tokens
            )  # Updated from parse_primary to parse_unary
            node = Node("comparison", op[1], [node, rhs])
        return node

    @staticmethod
    def parse_function_call(tokens):
        function_name = ExpressionParser.expect_token(tokens, "IDENTIFIER")[1]
        ExpressionParser.expect_token(tokens, "LPAREN")
        arguments = []
        while tokens[0][0] != "RPAREN":
            arguments.append(ExpressionParser.parse_expression(tokens))
            if tokens[0][0] != "RPAREN":
                ExpressionParser.expect_token(tokens, "COMMA")
        ExpressionParser.expect_token(tokens, "RPAREN")
        return Node("function_call", function_name, arguments)

    @staticmethod
    def parse_unary(tokens):
        if tokens and tokens[0][0] == "NEGATION":
            op = tokens.pop(0)
            operand = ExpressionParser.parse_unary(
                tokens
            )  # Unary negation can be nested
            return Node("unary_negation", op[1], [operand])
        else:
            return ExpressionParser.parse_primary(tokens)

    @staticmethod
    def parse_primary(tokens):
        if tokens[0][0] == "LPAREN":
            tokens.pop(0)
            node = ExpressionParser.parse_expression(tokens)
            ExpressionParser.expect_token(tokens, "RPAREN")
            return node
        elif tokens[0][0] == "IDENTIFIER":
            # Lookahead to check if this is a function call
            if len(tokens) > 1 and tokens[1][0] == "LPAREN":
                return ExpressionParser.parse_function_call(tokens)
            else:
                return Node("identifier", tokens.pop(0)[1])
        elif tokens[0][0] == "STRING":
            return Node("string", tokens.pop(0)[1])
        else:
            raise SyntaxError(f"Unexpected token: {tokens[0][1]}")

    def print_ast(self):
        self.__print_ast(self.parsed_expression)

    def get_node(self):
        return self.parsed_expression
