from __future__ import annotations

from typing import Optional

import lark
from starkware.cairo.lang.compiler.ast.expr import *
from starkware.cairo.lang.compiler.ast.expr import ExprIdentifier
from starkware.cairo.lang.compiler.error_handling import InputFile
from starkware.cairo.lang.compiler.parser_transformer import (
    ParserContext,
    ParserError,
    ParserTransformer,
)

import horus.compiler.parser
from horus.compiler.code_elements import (
    AnnotatedCodeElement,
    BoolConst,
    BoolExprAtom,
    BoolExprCompare,
    BoolNegation,
    BoolOperation,
    CodeElementAnnotation,
    CodeElementCheck,
    CodeElementLogicalVariableDeclaration,
    CodeElementStorageUpdate,
    ExprLogicalIdentifier,
)


class HorusTransformer(ParserTransformer):
    """
    Transforms lark trees generated by extended
    grammar `Horus.ebnf`.
    """

    def __init__(
        self,
        input_file: InputFile,
        parser_context: Optional[ParserContext],
        is_parsing_check: bool = False,
    ):
        self.is_parsing_check = is_parsing_check
        super().__init__(input_file, parser_context)

    @lark.v_args(meta=True)
    def commented_code_element(self, meta, value):
        comment: Optional[str] = (
            value[1][2:] if len(value) == 2 and value[1] is not None else None
        )

        if comment is not None:
            possible_annotation: str = comment.strip()
            if possible_annotation.startswith("@"):
                annotation_type = possible_annotation.split(" ")[0].strip()

                if annotation_type in [
                    "@pre",
                    "@post",
                    "@assert",
                    "@invariant",
                    "@declare",
                    "@storage_update",
                ]:
                    check = horus.compiler.parser.parse(
                        filename=self.input_file.filename,
                        code=possible_annotation,
                        code_type="annotation",
                        expected_type=CodeElementAnnotation,
                        parser_context=self.parser_context,
                    )

                    if isinstance(check, (CodeElementCheck, CodeElementStorageUpdate)):
                        check.unpreprocessed_rep = " ".join(
                            possible_annotation.split(" ")[1:]
                        )

                    code_elem = super().commented_code_element(meta, [value[0], None])
                    code_elem.code_elm = AnnotatedCodeElement(
                        check, code_elm=code_elem.code_elm
                    )
                    return code_elem
                else:
                    raise ParserError(
                        f"Unexpected annotation type: {annotation_type}",
                        location=self.meta2loc(meta),
                    )

        return super().commented_code_element(meta, value)

    @lark.v_args(meta=True)
    def logical_identifier(self, meta, value):
        identifier_name = ".".join(x.value for x in value)
        if value[0].value == "$Return":
            return ExprIdentifier(name=identifier_name, location=self.meta2loc(meta))
        else:
            return ExprLogicalIdentifier(
                name=identifier_name, location=self.meta2loc(meta)
            )

    @lark.v_args(meta=True)
    def logical_identifier_def(self, meta, value):
        if value[0].value == "$Return":
            raise ParserError(
                f"'$Return' is a reserved name.", location=self.meta2loc(meta)
            )
        return ExprLogicalIdentifier(name=value[0].value, location=self.meta2loc(meta))

    def atom_logical_identifier(self, value):
        return value[0]

    @lark.v_args(inline=True)
    def bool_formula_impl(self, lhs, rhs):
        return BoolOperation(lhs, rhs, "->")

    @lark.v_args(inline=True)
    def bool_formula_or(self, lhs, rhs):
        return BoolOperation(lhs, rhs, "|")

    @lark.v_args(inline=True)
    def bool_formula_and(self, lhs, rhs):
        return BoolOperation(lhs, rhs, "&")

    @lark.v_args(inline=True)
    def bool_unary_neg(self, formula):
        return BoolNegation(formula)

    @lark.v_args(inline=True)
    def bool_expr_true(self):
        return BoolConst(True)

    @lark.v_args(inline=True)
    def bool_expr_false(self):
        return BoolConst(False)

    @lark.v_args(inline=True)
    def bool_atom(self, expr):
        return BoolExprAtom(expr)

    @lark.v_args(inline=True, meta=True)
    def bool_expr_le(self, meta, lhs, rhs):
        return BoolExprCompare(lhs, rhs, "<=", location=self.meta2loc(meta))

    @lark.v_args(inline=True, meta=True)
    def bool_expr_lt(self, meta, lhs, rhs):
        return BoolExprCompare(lhs, rhs, "<", location=self.meta2loc(meta))

    @lark.v_args(inline=True, meta=True)
    def bool_expr_ge(self, meta, lhs, rhs):
        return BoolExprCompare(lhs, rhs, ">=", location=self.meta2loc(meta))

    @lark.v_args(inline=True, meta=True)
    def bool_expr_gt(self, meta, lhs, rhs):
        return BoolExprCompare(lhs, rhs, ">", location=self.meta2loc(meta))

    @lark.v_args(inline=True)
    def bool_expr_parentheses(self, formula):
        return formula

    @lark.v_args(inline=True, meta=True)
    def precond_annotation(self, meta, expr):
        return CodeElementCheck(
            CodeElementCheck.CheckKind.PRE_COND, expr, location=self.meta2loc(meta)
        )

    @lark.v_args(inline=True, meta=True)
    def postcond_annotation(self, meta, expr):
        return CodeElementCheck(
            CodeElementCheck.CheckKind.POST_COND, expr, location=self.meta2loc(meta)
        )

    @lark.v_args(inline=True, meta=True)
    def assert_annotation(self, meta, expr):
        return CodeElementCheck(
            CodeElementCheck.CheckKind.ASSERT, expr, location=self.meta2loc(meta)
        )

    @lark.v_args(inline=True, meta=True)
    def invariant_annotation(self, meta, expr):
        return CodeElementCheck(
            CodeElementCheck.CheckKind.INVARIANT, expr, location=self.meta2loc(meta)
        )

    @lark.v_args(inline=True, meta=True)
    def declare_annotation(self, meta, identifier, type):
        return CodeElementLogicalVariableDeclaration(
            identifier.name, type, location=self.meta2loc(meta)
        )

    @lark.v_args(inline=True, meta=True)
    def storage_update_annotation(self, meta, identifier, args, result):
        return CodeElementStorageUpdate(
            identifier.name, args, result, location=self.meta2loc(meta)
        )

    def transform(self, tree: lark.Tree):
        # The nodes of the tree imported from cairo.ebnf appear with
        # prefix starkware__cairo__lang__compiler__cairo__ which we remove here.
        for node in tree.iter_subtrees():
            node.data = node.data.replace(
                "starkware__cairo__lang__compiler__cairo__", ""
            )

        return super().transform(tree)
