from ASTTypes.ASTNodeSubtypes import CompoundStatement, Expression, StatementList
from ASTTypes.ASTNode import ContextType


class IfStatement(CompoundStatement):
    def __init__(self, condition: Expression,
                 block: StatementList,
                 else_block: StatementList | None = None):
        children = [condition, block]
        if else_block:
            children.append(else_block)
        super().__init__(children)

    @property
    def context(self) -> ContextType:
        return self._context

    @context.setter
    def context(self, value):
        self._context = value
        self.children[0].context = ContextType.SCALAR
        for child in self.children[1:]:
            child.context = ContextType.VOID

    @property
    def condition(self) -> Expression:
        return self.children[0]

    @property
    def block(self) -> Expression:
        return self.children[1]

    @property
    def else_block(self) -> Expression | None:
        try:
            return self.children[2]
        except IndexError:
            return None

    def append_else(self, else_block: StatementList):
        if len(self.children) < 3:
            self.children.append(else_block)
            return
        raise TypeError("Can not append another else statement!")

    def translate(self) -> str:
        block = "\t" + '\n\t'.join(self.block.translate().split('\n')[:-1]) + "\n"
        ret = f"if ({self.condition.translate()}) {{\n{block}}}"
        if self.else_block:
            if isinstance(self.else_block.children[0], self.__class__):
                ret += f" else {self.else_block.translate()}"
            else:
                block = "\t" + '\n\t'.join(self.else_block.translate().split('\n')[:-1]) + "\n"
                ret += f" else {{\n{block}}}"
        if ret[-1] != "\n":
            ret += "\n"
        return ret



class SubStatement(CompoundStatement):
    def __init__(self, name, block: StatementList):
        super().__init__([block])
        self.name = name

    def __repr__(self):
        return f"{super().__repr__()} {self.name}"

    def translate(self) -> str:
        block = "".join([child.translate() for child in self.children])
        block = "\t"+"\n\t".join(block.split('\n')[:-1])+"\n"
        return f"function {self.name} (..._){{\n{block}}}\n\n"


class ReturnStatement(CompoundStatement):
    def __init__(self, expression_to_return: Expression):
        super().__init__([expression_to_return])

    def translate(self) -> str:
        return f"return {self.children[0].translate()};\n"


class WhileStatement(CompoundStatement):
    def __init__(self, condition: Expression,
                 block: StatementList):
        children = [condition, block]
        super().__init__(children)

    @property
    def context(self) -> ContextType:
        return self._context

    @context.setter
    def context(self, value):
        self._context = value
        self.children[0].context = ContextType.SCALAR
        for child in self.children[1:]:
            child.context = ContextType.VOID

    @property
    def condition(self) -> Expression:
        return self.children[0]

    @property
    def block(self) -> Expression:
        return self.children[1]

    def translate(self) -> str:
        condition = self.condition.translate()
        block = self.block.translate()
        block = "\t" + "\n\t".join(block.split('\n')[:-1]) + "\n"
        return f"while ({condition}) {{\n{block}}}\n"
