from .tokens import Token
from .utils import *

OPS = "+-*/&|<>="
UNARY_OPS = "-~"
CONSTANTS = ["true", "false", "null", "this"]


class Term:
    def __init__(self):
        pass

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        if not tokens:
            return (None, 0)
        if tokens[0].type in ["integer", "string"] or tokens[0] in CONSTANTS:
            return (ConstantTerm(tokens[0]), 1)
        if tokens[0].token in UNARY_OPS:
            """Format: <unary op> <term>"""
            term, dt = Term.from_tokens(tokens[1:])
            return (UnaryTerm(tokens[0], term), dt + 1)
        if tokens[0] == LEFT_PAREN:
            """Format: '(' <expression> ')'"""
            expr, dt = Expression.from_tokens(tokens[1:])
            if tokens[dt + 1] != RIGHT_PAREN:
                raise UnaryTerm(RIGHT_PAREN, tokens[dt])
            return (expr, dt + 2)
        if tokens[0].type == "identifier":
            if tokens[1] in [LEFT_PAREN, "."]:
                # subroutine(...) or Class.subroutine(...)
                return SubroutineCall.from_tokens(tokens)
            if tokens[1] == LEFT_BRACKET:
                # array[index]
                return SubscriptTerm.from_tokens(tokens)
            return (VarTerm(tokens[0]), 1)
        return (None, 0)


class ConstantTerm:
    def __init__(self, term: Token):
        self.term = term

    def __str__(self):
        return self.term.token


class SubscriptTerm:
    def __init__(self, var: Token, subscript: Token):
        self.var = var
        self.subscript = subscript

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct a subscripted array term.

        Format:
        <var> '[' <expression> ']'
        """
        if tokens[0].type != "identifier" or tokens[1] != LEFT_BRACKET:
            return (None, 0)

        var = tokens[0]
        t = 2
        sub, dt = Expression.from_tokens(tokens[2:])
        if sub is None:
            raise JackSyntaxError(f"Expected subscript", tokens[2])
        t += dt
        if tokens[t] != RIGHT_BRACKET:
            raise UnexpectedToken(RIGHT_BRACKET, tokens[t])
        return (SubscriptTerm(var, sub), t + 1)

    def __str__(self):
        return f"{self.var}[{self.subscript}]"


class VarTerm:
    def __init__(self, var: Token):
        self.var = var

    def __str__(self):
        return self.var.token


class UnaryTerm:
    def __init__(self, op: Token, term: Term):
        self.op = op
        self.term = term

    def __str__(self):
        return f"({self.op}{self.term})"


class Expression:
    def __init__(self, lhs: Term, op=None, rhs=None):
        self.lhs = lhs
        self.op = op
        self.rhs = rhs

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

        Format:
        <term> (<op> <term>)?
        """
        lhs, dt = Term.from_tokens(tokens)
        if lhs is None:
            return (None, 0)
        t = dt

        op = tokens[t]
        if op.token not in OPS:
            return (Expression(lhs), t)

        t += 1
        rhs, dt = Term.from_tokens(tokens[t:])
        if rhs is None:
            raise UnexpectedToken("other term", rhs)
        t += dt

        return (Expression(lhs, op, rhs), t)

    def __str__(self):
        if self.op is not None:
            return f"({self.lhs} {self.op} {self.rhs})"
        else:
            return str(self.lhs)


class ExpressionList:
    def __init__(self, exprs: list[Expression]):
        self.exprs = exprs

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

        Format:
        (<expression> (',' <expression>)*)?
        """
        t = 0
        exprs = []
        while True:
            expr, dt = Expression.from_tokens(tokens[t:])
            if expr is None:
                if t == 0:
                    # only allow lack of expression right after paren
                    break
                else:
                    # expect expression after comma
                    raise JackSyntaxError(f"Expected expression", tokens[t])
            t += dt
            exprs.append(expr)
            if tokens[t] != ",":
                break
            t += 1

        return (ExpressionList(exprs), t)

    def __str__(self):
        return ", ".join([str(expr) for expr in self.exprs])


class SubroutineCall:
    def __init__(self, jack_class: Token, name: Token, exprs: ExpressionList):
        self.jack_class = jack_class
        self.name = name
        self.exprs = exprs

    @classmethod
    def from_tokens(cls, tokens: list) -> tuple:
        """Construct invocation of subroutine.

        Format:
        (<class> '.')? <subroutine> '(' <expression list> ')'
        """
        t = 0
        jack_class = None
        if tokens[1] == ".":
            jack_class = tokens[0]
            t = 2

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

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

        exprs, dt = ExpressionList.from_tokens(tokens[t:])
        t += dt
        if tokens[t] != RIGHT_PAREN:
            raise UnexpectedToken(RIGHT_PAREN, tokens[t])
        return (SubroutineCall(jack_class, name, exprs), t + 1)

    def __str__(self):
        if self.jack_class is None:
            return f"{self.name}({self.exprs})"
        return f"{self.jack_class}.{self.name}({self.exprs})"