from typing import List, Optional
from .lexer import Token, TokenType, VerilogLexer
from .ast_nodes import (
    Module, Port, WireDeclaration, RegDeclaration, Assign, Always,
    Identifier, Number, BinaryOp, UnaryOp, Statement, Declaration, Expression
)


class VerilogParser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def error(self, message: str):
        if self.pos < len(self.tokens):
            token = self.tokens[self.pos]
            raise SyntaxError(f"Parser error at line {token.line}, column {token.column}: {message}")
        else:
            raise SyntaxError(f"Parser error at end of file: {message}")

    def peek(self, offset: int = 0) -> Optional[Token]:
        pos = self.pos + offset
        return self.tokens[pos] if pos < len(self.tokens) else None

    def advance(self) -> Token:
        if self.pos < len(self.tokens):
            token = self.tokens[self.pos]
            self.pos += 1
            return token
        self.error("Unexpected end of file")

    def expect(self, token_type: TokenType, value: Optional[str] = None) -> Token:
        token = self.peek()
        if not token or token.type != token_type:
            self.error(f"Expected {token_type}")
        if value is not None and token.value != value:
            self.error(f"Expected '{value}'")
        return self.advance()

    def match(self, token_type: TokenType, value: Optional[str] = None) -> bool:
        token = self.peek()
        if not token:
            return False
        if token.type != token_type:
            return False
        if value is not None and token.value != value:
            return False
        return True

    def consume(self, token_type: TokenType, value: Optional[str] = None) -> bool:
        if self.match(token_type, value):
            self.advance()
            return True
        return False

    def parse(self) -> Module:
        return self.parse_module()

    def parse_module(self) -> Module:
        self.expect(TokenType.KEYWORD, 'module')
        name_token = self.expect(TokenType.IDENTIFIER)
        module_name = name_token.value

        port_list = []
        if self.consume(TokenType.DELIMITER, '('):
            port_list = self.parse_port_list()
            self.expect(TokenType.DELIMITER, ')')

        self.expect(TokenType.DELIMITER, ';')

        declarations = []
        statements = []

        while not self.match(TokenType.KEYWORD, 'endmodule'):
            if self.match(TokenType.KEYWORD, 'wire'):
                declarations.append(self.parse_wire_declaration())
            elif self.match(TokenType.KEYWORD, 'reg'):
                declarations.append(self.parse_reg_declaration())
            elif self.match(TokenType.KEYWORD, 'assign'):
                statements.append(self.parse_assign())
            elif self.match(TokenType.KEYWORD, 'always'):
                statements.append(self.parse_always())
            elif self.match(TokenType.KEYWORD, 'input') or self.match(TokenType.KEYWORD, 'output'):
                declarations.append(self.parse_io_declaration())
            else:
                self.advance()

        self.expect(TokenType.KEYWORD, 'endmodule')

        return Module(name=module_name, port_list=port_list, declarations=declarations, statements=statements)

    def parse_port_list(self) -> List[Port]:
        ports = []
        while not self.match(TokenType.DELIMITER, ')'):
            port = self.parse_port()
            ports.append(port)
            if not self.consume(TokenType.DELIMITER, ','):
                break
        return ports

    def parse_port(self) -> Port:
        direction_token = self.advance()
        direction = direction_token.value

        width = None
        if self.consume(TokenType.DELIMITER, '['):
            width_start = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ':')
            width_end = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ']')
            width = (int(width_start), int(width_end))

        name_token = self.expect(TokenType.IDENTIFIER)
        return Port(name=name_token.value, direction=direction, width=width)

    def parse_wire_declaration(self) -> WireDeclaration:
        self.expect(TokenType.KEYWORD, 'wire')

        width = None
        if self.consume(TokenType.DELIMITER, '['):
            start = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ':')
            end = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ']')
            width = (int(start), int(end))

        name_token = self.expect(TokenType.IDENTIFIER)
        self.expect(TokenType.DELIMITER, ';')

        return WireDeclaration(name=name_token.value, width=width)

    def parse_reg_declaration(self) -> RegDeclaration:
        self.expect(TokenType.KEYWORD, 'reg')

        width = None
        if self.consume(TokenType.DELIMITER, '['):
            start = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ':')
            end = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ']')
            width = (int(start), int(end))

        name_token = self.expect(TokenType.IDENTIFIER)
        self.expect(TokenType.DELIMITER, ';')

        return RegDeclaration(name=name_token.value, width=width)

    def parse_io_declaration(self) -> Declaration:
        direction_token = self.advance()
        width = None
        if self.consume(TokenType.DELIMITER, '['):
            start = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ':')
            end = self.expect(TokenType.NUMBER).value
            self.expect(TokenType.DELIMITER, ']')
            width = (int(start), int(end))

        name_token = self.expect(TokenType.IDENTIFIER)
        self.expect(TokenType.DELIMITER, ';')

        if direction_token.value == 'input':
            return WireDeclaration(name=name_token.value, width=width)
        else:
            return RegDeclaration(name=name_token.value, width=width)

    def parse_assign(self) -> Assign:
        self.expect(TokenType.KEYWORD, 'assign')
        left = self.parse_expression()
        if self.consume(TokenType.DELIMITER, '='):
            pass
        elif self.consume(TokenType.OPERATOR, '='):
            pass
        else:
            self.error("Expected '='")
        right = self.parse_expression()
        self.expect(TokenType.DELIMITER, ';')
        return Assign(left=left, right=right)

    def parse_always(self) -> Always:
        self.expect(TokenType.KEYWORD, 'always')
        self.expect(TokenType.DELIMITER, '@')
        self.expect(TokenType.DELIMITER, '(')
        sensitivity_list = self.parse_sensitivity_list()
        self.expect(TokenType.DELIMITER, ')')

        body = []
        if self.consume(TokenType.KEYWORD, 'begin'):
            while not self.match(TokenType.KEYWORD, 'end'):
                statement = self.parse_statement()
                if statement:
                    body.append(statement)
                if self.peek() and self.peek().type == TokenType.DELIMITER and self.peek().value == ';':
                    self.advance()
            self.expect(TokenType.KEYWORD, 'end')
        else:
            statement = self.parse_statement()
            if statement:
                body.append(statement)

        return Always(sensitivity_list=sensitivity_list, body=body)

    def parse_sensitivity_list(self) -> List[str]:
        signals = []
        while self.peek() and not self.match(TokenType.DELIMITER, ')'):
            if self.match(TokenType.KEYWORD, 'posedge') or self.match(TokenType.KEYWORD, 'negedge'):
                edge_token = self.advance()
                if self.match(TokenType.IDENTIFIER):
                    signal_token = self.advance()
                    signals.append(f"{edge_token.value} {signal_token.value}")
            elif self.match(TokenType.IDENTIFIER):
                signal_token = self.advance()
                signals.append(signal_token.value)
            elif self.consume(TokenType.KEYWORD, 'or'):
                continue
            else:
                break
        return signals

    def parse_statement(self) -> Optional[Statement]:
        expr = self.parse_expression()
        if self.consume(TokenType.DELIMITER, '=') or self.consume(TokenType.OPERATOR, '<='):
            right_expr = self.parse_expression()
            return Assign(left=expr, right=right_expr)
        return None

    def parse_expression(self) -> Expression:
        return self.parse_or_expression()

    def parse_or_expression(self) -> Expression:
        left = self.parse_and_expression()
        while self.peek() and self.peek().value in ['||', 'or']:
            op_token = self.advance()
            right = self.parse_and_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_and_expression(self) -> Expression:
        left = self.parse_equality_expression()
        while self.peek() and self.peek().value in ['&&', 'and', '&']:
            op_token = self.advance()
            right = self.parse_equality_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_equality_expression(self) -> Expression:
        left = self.parse_relational_expression()
        while self.peek() and self.peek().value in ['==', '!=']:
            op_token = self.advance()
            right = self.parse_relational_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_relational_expression(self) -> Expression:
        left = self.parse_shift_expression()
        while self.peek() and self.peek().value in ['<', '>', '<=', '>=']:
            op_token = self.advance()
            right = self.parse_shift_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_shift_expression(self) -> Expression:
        left = self.parse_additive_expression()
        while self.peek() and self.peek().value in ['<<', '>>']:
            op_token = self.advance()
            right = self.parse_additive_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_additive_expression(self) -> Expression:
        left = self.parse_multiplicative_expression()
        while self.peek() and self.peek().value in ['+', '-']:
            op_token = self.advance()
            right = self.parse_multiplicative_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_multiplicative_expression(self) -> Expression:
        left = self.parse_unary_expression()
        while self.peek() and self.peek().value in ['*', '/', '%']:
            op_token = self.advance()
            right = self.parse_unary_expression()
            left = BinaryOp(left=left, operator=op_token.value, right=right)
        return left

    def parse_unary_expression(self) -> Expression:
        if self.peek() and self.peek().value in ['!', '~', '-', '+']:
            op_token = self.advance()
            operand = self.parse_unary_expression()
            return UnaryOp(operator=op_token.value, operand=operand)
        return self.parse_primary_expression()

    def parse_primary_expression(self) -> Expression:
        if self.match(TokenType.NUMBER):
            token = self.advance()
            return Number(value=token.value)
        elif self.match(TokenType.IDENTIFIER):
            token = self.advance()
            return Identifier(name=token.value)
        elif self.consume(TokenType.DELIMITER, '('):
            expr = self.parse_expression()
            self.expect(TokenType.DELIMITER, ')')
            return expr
        else:
            self.error("Expected expression")
