from .expressions import Expression, SubroutineCall
from .utils import *


class Statement:
    def __init__(self):
        pass

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        for StatementClass in [
            LetStatement,
            DoStatement,
            ReturnStatement,
            IfStatement,
            WhileStatement,
        ]:
            stmt, dt = StatementClass.from_tokens(tokens)
            if stmt is not None:
                return (stmt, dt)

        return (None, 0)


class StatementList:
    def __init__(self, statements: list[Statement]):
        self.statements = statements

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        t = 0
        statements = []
        while True:
            stmt, dt = Statement.from_tokens(tokens[t:])
            if stmt is None:
                break
            statements.append(stmt)
            t += dt
        return (StatementList(statements), t)

    def print_verbose(self):
        for stmt in self.statements:
            stmt.print_verbose()
        if not self.statements:
            print("Do nothing")


class LetStatement:
    def __init__(self, name, expr):
        self.name = name
        self.expr = expr

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct let statement.

        Format:
        'let' <name> '=' <expression> ';'
        """
        if len(tokens) < 5 or tokens[0] != "let":
            return (None, 0)

        name = tokens[1]
        if name.type != "identifier":
            raise UnexpectedToken("variable name", name)

        if tokens[2] != "=":
            raise UnexpectedToken("=", tokens[2])

        t = 3
        expr, dt = Expression.from_tokens(tokens[t:])
        if expr is None:
            raise UnexpectedToken(f"Expected expression", tokens[3])
        t += dt

        if tokens[t] != ";":
            raise UnexpectedToken(";", tokens[t])

        return (LetStatement(name, expr), t + 1)

    def print_verbose(self):
        print(f"Let {self.name} be {self.expr}")


class DoStatement:
    def __init__(self, subcall: SubroutineCall):
        self.subcall = subcall

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct do statement.

        Format:
        'do' <subroutine call> ';'
        """
        if len(tokens) < 5 or tokens[0] != "do":
            return (None, 0)

        t = 1
        subcall, dt = SubroutineCall.from_tokens(tokens[1:])
        if subcall is None:
            raise UnexpectedToken("subroutine call", tokens[1])
        t += dt

        if tokens[t] != ";":
            raise UnexpectedToken(";", tokens[t])

        return (DoStatement(subcall), t + 1)

    def print_verbose(self):
        print(f"Do {self.subcall}")


class ReturnStatement:
    def __init__(self, expr: Expression):
        self.expr = expr

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct return statement.

        Format:
        'return' <expression>? ';'
        """
        if len(tokens) < 3 or tokens[0] != "return":
            return (None, 0)

        t = 1
        expr, dt = Expression.from_tokens(tokens[1:])
        t += dt

        if tokens[t] != ";":
            raise UnexpectedToken(";", tokens[t])

        return (ReturnStatement(expr), t + 1)

    def print_verbose(self):
        print(f"Return {self.expr or 'null'}")


class IfStatement:
    def __init__(self, condition: Expression, then: StatementList, else_then=None):
        self.condition = condition
        self.then = then
        self.else_then = else_then

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct if statement.

        Format:
        'if' '(' <expression> ')' '{' <statement list> '}' ('else' '{' <statement list '}')?
        """
        if len(tokens) < 6 or tokens[0] != "if":
            return (None, 0)

        if tokens[1] != LEFT_PAREN:
            raise UnexpectedToken(LEFT_PAREN, tokens[1])

        t = 2
        cond, dt = Expression.from_tokens(tokens[2:])
        if cond is None:
            raise JackSyntaxError(f"Expected condition", tokens[2])
        t += dt

        if tokens[t] != RIGHT_PAREN:
            raise UnexpectedToken(RIGHT_PAREN, tokens[t])
        if tokens[t + 1] != LEFT_BRACE:
            raise UnexpectedToken(LEFT_BRACE, tokens[t + 1])
        t += 2

        then, dt = StatementList.from_tokens(tokens[t:])
        t += dt

        if tokens[t] != RIGHT_BRACE:
            raise UnexpectedToken(RIGHT_BRACE, tokens[t])
        t += 1

        if tokens[t] != "else":
            return (IfStatement(cond, then), t)

        if tokens[t + 1] != LEFT_BRACE:
            raise UnexpectedToken(LEFT_BRACE, tokens[t + 1])

        t += 2

        else_then, dt = StatementList.from_tokens(tokens[t:])
        t += dt

        if tokens[t] != RIGHT_BRACE:
            raise UnexpectedToken(RIGHT_BRACE, tokens[t])

        return (IfStatement(cond, then, else_then), t + 1)

    def print_verbose(self):
        print(f"If {self.condition} then:")
        self.then.print_verbose()
        if self.else_then is not None:
            print("Else then:")
            self.else_then.print_verbose()
        print("End if")


class WhileStatement:
    def __init__(self, condition: Expression, do: StatementList):
        self.condition = condition
        self.do = do

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct while statement.

        Format:
        'while' '(' <expression> ')' '{' <statement list> '}'
        """
        if len(tokens) < 6 or tokens[0] != "while":
            return (None, 0)

        if tokens[1] != LEFT_PAREN:
            raise UnexpectedToken(LEFT_PAREN, tokens[1])

        t = 2
        cond, dt = Expression.from_tokens(tokens[2:])
        if cond is None:
            raise JackSyntaxError(f"Expected condition", tokens[2])
        t += dt

        if tokens[t] != RIGHT_PAREN:
            raise UnexpectedToken(RIGHT_PAREN, tokens[t])
        if tokens[t + 1] != LEFT_BRACE:
            raise UnexpectedToken(LEFT_BRACE, tokens[t + 1])
        t += 2

        do, dt = StatementList.from_tokens(tokens[t:])
        t += dt

        if tokens[t] != RIGHT_BRACE:
            raise UnexpectedToken(RIGHT_BRACE, tokens[t])

        return (WhileStatement(cond, do), t + 1)

    def print_verbose(self):
        print(f"While {self.condition} do:")
        self.do.print_verbose()
        print("End while")