from typing import Text, List, Dict, Any

import pytest

import rasa.shared.nlu.training_data.entities_parser as entities_parser
from rasa.shared.nlu.constants import TEXT


@pytest.mark.parametrize(
    "example, expected_entities, expected_text",
    [
        (
            "I need an [economy class](travel_flight_class:economy) ticket from "
            '[boston]{"entity": "city", "role": "from"} to [new york]{"entity": "city",'
            ' "role": "to"}, please.',
            [
                {
                    "start": 10,
                    "end": 23,
                    "value": "economy",
                    "entity": "travel_flight_class",
                },
                {
                    "start": 36,
                    "end": 42,
                    "value": "boston",
                    "entity": "city",
                    "role": "from",
                },
                {
                    "start": 46,
                    "end": 54,
                    "value": "new york",
                    "entity": "city",
                    "role": "to",
                },
            ],
            "I need an economy class ticket from boston to new york, please.",
        ),
        ("i'm looking for a place to eat", [], "i'm looking for a place to eat"),
        (
            "i'm looking for a place in the [north](loc-direction) of town",
            [{"start": 31, "end": 36, "value": "north", "entity": "loc-direction"}],
            "i'm looking for a place in the north of town",
        ),
        (
            "show me [chines](cuisine:chinese) restaurants",
            [{"start": 8, "end": 14, "value": "chinese", "entity": "cuisine"}],
            "show me chines restaurants",
        ),
        (
            'show me [italian]{"entity": "cuisine", "value": "22_ab-34*3.A:43er*+?df"} '
            "restaurants",
            [
                {
                    "start": 8,
                    "end": 15,
                    "value": "22_ab-34*3.A:43er*+?df",
                    "entity": "cuisine",
                }
            ],
            "show me italian restaurants",
        ),
        ("Do you know {ABC} club?", [], "Do you know {ABC} club?"),
        (
            "show me [chines](22_ab-34*3.A:43er*+?df) restaurants",
            [{"start": 8, "end": 14, "value": "43er*+?df", "entity": "22_ab-34*3.A"}],
            "show me chines restaurants",
        ),
        (
            'I want to fly from [Berlin]{"entity": "city", "role": "to"} to [LA]{'
            '"entity": "city", "role": "from", "value": "Los Angeles"}',
            [
                {
                    "start": 19,
                    "end": 25,
                    "value": "Berlin",
                    "entity": "city",
                    "role": "to",
                },
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "city",
                    "role": "from",
                },
            ],
            "I want to fly from Berlin to LA",
        ),
        (
            'I want to fly from [Berlin](city) to [LA]{"entity": "city", "role": '
            '"from", "value": "Los Angeles"}',
            [
                {"start": 19, "end": 25, "value": "Berlin", "entity": "city"},
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "city",
                    "role": "from",
                },
            ],
            "I want to fly from Berlin to LA",
        ),
    ],
)
def test_markdown_entity_regex(
    example: Text, expected_entities: List[Dict[Text, Any]], expected_text: Text
):

    result = entities_parser.find_entities_in_training_example(example)
    assert result == expected_entities

    replaced_text = entities_parser.replace_entities(example)
    assert replaced_text == expected_text


def test_parse_training_example():
    message = entities_parser.parse_training_example("Hello!", intent="greet")
    assert message.get("intent") == "greet"
    assert message.get(TEXT) == "Hello!"


def test_parse_empty_example():
    message = entities_parser.parse_training_example("")
    assert message.get("intent") is None
    assert message.get(TEXT) == ""


def test_parse_training_example_with_entities():
    message = entities_parser.parse_training_example(
        "I am from [Berlin](city).", intent="inform"
    )
    assert message.get("intent") == "inform"
    assert message.get(TEXT) == "I am from Berlin."
    assert message.get("entities") == [
        {"start": 10, "end": 16, "value": "Berlin", "entity": "city"}
    ]

@@ -122,16 +122,80 @@
 _parser%0A
+from rasa.shared.exceptions import InvalidEntityFormatException%0A
 from ras
@@ -4788,8 +4788,320 @@
 %7D%0A    %5D%0A
+%0A%0Adef test_markdown_entity_regex_error_handling():%0A%0A    with pytest.raises(InvalidEntityFormatException):%0A        entities_parser.find_entities_in_training_example(%0A            # JSON syntax error: missing closing %22 for %60role%60%0A            'I want to fly from %5BBerlin%5D%7B%22entity%22: %22city%22, %22role: %22from%22%7D'%0A        )%0A

add test for error handling in find_entities_in_training_example()