compiler.py

class ErasCompiler:
    def __init__(self):
        self.KEYWORDS = {
            "ARE YOU READY FOR IT?": "BEGIN_MAIN",
            "LONG LIVE": "END_MAIN",
            "BLANK SPACE": "DECLARE_VAR",
            "YOU BELONG WITH ME": "ASSIGN_VAL",
            "SHAKE IT OFF": "BEGIN_MATH",
            "MINE": "ADD",
            "BACK TO DECEMBER": "SUBTRACT",
            "CALL IT WHAT YOU WANT": "END_MATH",
            "SPEAK NOW": "PRINT",
            "I KNEW YOU WERE TROUBLE": "IF",
            "YOU'RE ON YOUR OWN, KID": "ELSE",
            "EXILE": "END_IF",
            "STAY STAY STAY": "==",
            "LET'S GO BATTLE": ">",
            "FROM THE VAULT": "BEGIN_FUNC",
            "THE STORY OF US": "FUNC_ARG",
            "CLEAN": "END_FUNC",
            "IMGONNAGETYOUBACK": "RETURN",
            "ME!": "AND",
            "THE 1": "OR",
            "LOOK WHAT YOU MADE ME DO": "NOT",
            "QUESTION...?": "INPUT",
            "WE ARE NEVER EVER GETTING BACK TOGETHER": "BREAK",
            "IS IT OVER NOW?": "WHILE",
            "OUT OF THE WOODS": "END_WHILE",
        }

    def transpile(self, source_code: str) -> str:
        lines = source_code.splitlines()
        python_output = ["enchanted = 1\nexile = 0\nthirteen = 13\n"]
        indent_level = 0
        math_target = None
        last_var = None
        in_function_def = False  # Track if we're building a function definition
        func_arg_count = 0  # Count function arguments

        for line in lines:
            raw = line  # Save original line for counting
            line = line.strip()
            if not line or line.startswith("DEAR JOHN"):
                continue

            token = None
            payload = ""
            for lyric, t in self.KEYWORDS.items():
                if line.startswith(lyric):
                    token = t
                    payload = line.replace(lyric, "").strip()
                    break
            
            if not token:
                continue

            # If we're in a function definition and encounter a non-FUNC_ARG/END_FUNC token,
            # close the function definition first
            if in_function_def and token not in ["FUNC_ARG", "END_FUNC"]:
                if python_output and not python_output[-1].endswith(":\n"):
                    if python_output[-1].endswith("("):
                        python_output[-1] += "):\n"
                    else:
                        python_output[-1] += "):\n"
                    indent_level += 1  # Increment for function body
                in_function_def = False
                func_arg_count = 0

            # Replace comparison operators in payload
            payload = payload.replace("STAY STAY STAY", "==")
            payload = payload.replace("LET'S GO BATTLE", ">")
            # Replace input keyword in payload - this handles "YOU BELONG WITH ME QUESTION...?"
            payload = payload.replace("QUESTION...?", "int(input())")
            # Remove PLAY keyword (used for function calls)
            payload = payload.replace("PLAY ", "")

            current_indent = "    " * indent_level

            # Core Logic Mapping
            if token == "BEGIN_MAIN":
                python_output.append("def main():\n")
                indent_level += 1
            elif token == "END_MAIN":
                python_output.append("\nmain()\n")  # Simplified for API execution
            elif token == "DECLARE_VAR":
                last_var = payload
                python_output.append(f"{current_indent}{payload} = 0\n")
            elif token == "ASSIGN_VAL":
                # payload already has "int(input())" if it contained "QUESTION...?"
                val = payload
                if math_target:
                    python_output.append(f"{current_indent}_acc = {val}\n")
                else:
                    python_output.append(f"{current_indent}{last_var} = {val}\n")
            elif token == "BEGIN_MATH":
                math_target = payload
                # Initialize accumulator with the target variable
                python_output.append(f"{current_indent}_acc = {payload}\n")
            elif token == "END_MATH":
                # Use payload if provided, otherwise use math_target
                target_var = payload if payload else math_target
                if target_var:
                    python_output.append(f"{current_indent}{target_var} = _acc\n")
                math_target = None
            elif token == "ADD":
                python_output.append(f"{current_indent}_acc += {payload}\n")
            elif token == "SUBTRACT":
                python_output.append(f"{current_indent}_acc -= {payload}\n")
            elif token == "AND":
                python_output.append(f"{current_indent}_acc = int(bool(_acc) and bool({payload}))\n")
            elif token == "OR":
                python_output.append(f"{current_indent}_acc = int(bool(_acc) or bool({payload}))\n")
            elif token == "NOT":
                # Count how many times Taylor told you to look
                flip_count = raw.count("LOOK WHAT YOU MADE ME DO")
                
                # We start with the value (either the payload variable or the current accumulator)
                target = payload if payload else "_acc"
                
                # Use modulo 2 instead of spamming nots
                if flip_count % 2 == 1:
                    python_output.append(f"{current_indent}_acc = int(not bool({target}))\n")
                else:
                    python_output.append(f"{current_indent}_acc = int(bool({target}))\n")
            elif token == "==":
                python_output.append(f"{current_indent}_acc = int(_acc == {payload})\n")
            elif token == "WHILE":
                python_output.append(f"{current_indent}while {payload}:\n")
                indent_level += 1
            elif token == "END_WHILE":
                indent_level -= 1
            elif token == "IF":
                python_output.append(f"{current_indent}if {payload}:\n")
                indent_level += 1
            elif token == "ELSE":
                indent_level -= 1
                python_output.append(f"{'    ' * indent_level}else:\n")
                indent_level += 1
            elif token == "END_IF":
                indent_level -= 1
            elif token == "PRINT":
                python_output.append(f"{current_indent}print({payload})\n")
            elif token == "BREAK":
                python_output.append(f"{current_indent}break\n")
            elif token == "INPUT":
                python_output.append(f"{current_indent}{last_var} = int(input())\n")
            elif token == "BEGIN_FUNC":
                python_output.append(f"{current_indent}def {payload}(")
                in_function_def = True
                func_arg_count = 0
            elif token == "FUNC_ARG":
                func_arg_count += 1
                # Check if this is the first argument (last line ends with "(")
                if python_output and python_output[-1].endswith("("):
                    python_output[-1] += payload
                else:
                    # Add comma and argument
                    python_output[-1] += f", {payload}"
            elif token == "END_FUNC":
                # Close the function definition if it's still open
                if in_function_def and python_output and not python_output[-1].endswith(":\n"):
                    if python_output[-1].endswith("("):
                        # No arguments, close with ):
                        python_output[-1] += "):\n"
                    else:
                        # Has arguments, close with ):
                        python_output[-1] += "):\n"
                in_function_def = False
                func_arg_count = 0
                indent_level -= 1
            elif token == "RETURN":
                python_output.append(f"{current_indent}return {payload}\n")

        return "".join(python_output)