#!/usr/bin/env python
"""
This is an example of "prompt_toolkit.contrib.regular_languages" which
implements a little calculator.

Type for instance::

    > add 4 4
    > sub 4 4
    > sin 3.14

This example shows how you can define the grammar of a regular language and how
to use variables in this grammar with completers and tokens attached.
"""

import math

from prompt_toolkit import prompt
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.contrib.regular_languages.compiler import compile
from prompt_toolkit.contrib.regular_languages.completion import GrammarCompleter
from prompt_toolkit.contrib.regular_languages.lexer import GrammarLexer
from prompt_toolkit.lexers import SimpleLexer
from prompt_toolkit.styles import Style

operators1 = ["add", "sub", "div", "mul"]
operators2 = ["cos", "sin"]


def create_grammar():
    return compile(
        r"""
        (\s*  (?P<operator1>[a-z]+)   \s+   (?P<var1>[0-9.]+)   \s+   (?P<var2>[0-9.]+)   \s*) |
        (\s*  (?P<operator2>[a-z]+)   \s+   (?P<var1>[0-9.]+)   \s*)
    """
    )


example_style = Style.from_dict(
    {
        "operator": "#33aa33 bold",
        "number": "#ff0000 bold",
        "trailing-input": "bg:#662222 #ffffff",
    }
)


if __name__ == "__main__":
    g = create_grammar()

    lexer = GrammarLexer(
        g,
        lexers={
            "operator1": SimpleLexer("class:operator"),
            "operator2": SimpleLexer("class:operator"),
            "var1": SimpleLexer("class:number"),
            "var2": SimpleLexer("class:number"),
        },
    )

    completer = GrammarCompleter(
        g,
        {
            "operator1": WordCompleter(operators1),
            "operator2": WordCompleter(operators2),
        },
    )

    try:
        # REPL loop.
        while True:
            # Read input and parse the result.
            text = prompt(
                "Calculate: ", lexer=lexer, completer=completer, style=example_style
            )
            m = g.match(text)
            if m:
                vars = m.variables()
            else:
                print("Invalid command\n")
                continue

            print(vars)
            if vars.get("operator1") or vars.get("operator2"):
                try:
                    var1 = float(vars.get("var1", 0))
                    var2 = float(vars.get("var2", 0))
                except ValueError:
                    print("Invalid command (2)\n")
                    continue

                # Turn the operator string into a function.
                operator = {
                    "add": (lambda a, b: a + b),
                    "sub": (lambda a, b: a - b),
                    "mul": (lambda a, b: a * b),
                    "div": (lambda a, b: a / b),
                    "sin": (lambda a, b: math.sin(a)),
                    "cos": (lambda a, b: math.cos(a)),
                }[vars.get("operator1") or vars.get("operator2")]

                # Execute and print the result.
                print(f"Result: {operator(var1, var2)}\n")

            elif vars.get("operator2"):
                print("Operator 2")

    except EOFError:
        pass