import timeit
from typing import Callable, Optional

import pytest
from django.conf import settings

from promptgame.gameui import restrictions
from promptgame.gameui.restrictions import RESTRICTIONS_DICT, RestrictionResult


def helper_test_restriction(
    restriction_id: Optional[str],
    text: str,
    expected: bool,
    partial_error_message: Optional[str] = None,
    restriction_fn: Optional[Callable[[str], RestrictionResult]] = None,
):
    if restriction_fn is None:
        restriction = RESTRICTIONS_DICT[restriction_id]
        restriction_fn = restriction.is_valid
    result = restriction_fn(text)
    assert result.is_valid == expected
    if partial_error_message:
        assert partial_error_message in result.error_message
    else:
        assert not result.error_message


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("This is a simple test.", True, None),
        (" hi" * 51, False, "more than 50"),  # Too long
        (" hi" * 50, True, None),  # hi is a common word, so should be a single token
    ],
)
def test_token_limit_50(text, expected, partial_error_message):
    helper_test_restriction("token_limit_50", text, expected, partial_error_message)


COMMON_ERR = "These words are not in the 1,000 most common"
ALPHA_SPACE_ERR = restrictions.CHAR_SET_DICT["alpha_space"].message


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("the of and to", True, None),  # common words
        ("THE OF AND TO", True, None),  # common words capitalized
        ("the. of and to", False, ALPHA_SPACE_ERR),  # common words with punctuation
        ("the \t to \n and  ", True, None),  # Tabs, newlines, multi spaces are OK
        ("the don't and to", False, ALPHA_SPACE_ERR),  # common words with punctuation
        ("the of 'and' to", False, ALPHA_SPACE_ERR),  # common words with punctuation
        (
            "the of.and to",
            False,
            ALPHA_SPACE_ERR,
        ),  # common words joined with punctuation
        ("octagonal seismic rhetorical", False, COMMON_ERR),  # less common words
        ("octagonal the and", False, COMMON_ERR),  # only one uncommon word
        ("o c t a g o n a l", False, COMMON_ERR),  # single chars aren't common
        ("I i", True, None),  # I is the only capitalized common word
    ],
)
def test_common_1k_words(text, expected, partial_error_message):
    helper_test_restriction("common_1k_words", text, expected, partial_error_message)


UNCOMMON_ERR = "in the 1,000 most common"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("the of and to", False, UNCOMMON_ERR),  # common words
        ("THE OF AND TO", False, UNCOMMON_ERR),  # common words capitalized
        ("octagonal.", True, None),  # uncommon words with punctuation
        ("octagonal \t seismic \n rhetorical  ", True, None),  # whitespace OK
        ("octagonal seismic rhetorical", True, None),  # less common words
        ("octagonal the and", False, UNCOMMON_ERR),  # only one uncommon word
        ("octasdfasdfasf", True, None),  # fake words OK
    ],
)
def test_not_common_1k_words(text, expected, partial_error_message):
    helper_test_restriction(
        "not_common_1k_words", text, expected, partial_error_message
    )


NOT_ENGLISH_ERR = "not valid English words"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("the THE The tHe", True, None),  # OK with any capitalization
        ("the.", True, None),  # okay with punctuation
        ("'the'", True, None),  # okay with punctuation
        ("asdf", False, NOT_ENGLISH_ERR),  # error on non-words
        ("principio", False, NOT_ENGLISH_ERR),  # error on non-English words
        ("tiburón", False, NOT_ENGLISH_ERR),  # error on non-English words
        ("你好", False, NOT_ENGLISH_ERR),  # error on words in other alphabets
        ("𝕋𝕖𝕩𝕥 𝔼𝕕𝕚𝕥𝕠𝕣", False, NOT_ENGLISH_ERR),  # error on special characters
        ("🌕", True, None),  # Emojis are fine
        ("this.is.a.test", True, None),  # punctuation splits words
        ("color", True, None),  # OK with American spelling
        ("colour", True, None),  # OK with British spelling
        ("the \t unit \n test  ", True, None),  # whitespace OK
        ("the     unit \n\n test  ", True, None),  # many spaces OK
        ("ooo and vvv", False, "ooo, vvv"),  # words mentioned in order
        ("vvv and ooo", False, "vvv, ooo"),  # words mentioned in order
    ],
)
def test_check_valid_english_words(text, expected, partial_error_message):
    helper_test_restriction(
        None,
        text,
        expected,
        partial_error_message,
        restrictions.check_valid_english_words,
    )


ALNUM_SPACE_PUNCT_ERR = restrictions.CHAR_SET_DICT["alnum_space_punct"].message


BANNED_ERR = "Using banned"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("yolo 123", True, None),  # OK with words that violate no restrictions
        (
            "hell the",
            True,
            None,
        ),  # OK with words that contain substrings of banned words
        ("hello world", False, BANNED_ERR),  # Error on banned word
        ("Hello world", False, BANNED_ERR),  # Error on capitalized word
        ("bad", False, BANNED_ERR),  # Error on banned word as part of other word
        ("for the win", False, BANNED_ERR),  # Error on multi-word banned word
        ("say hello for the win", False, "hello"),  # Mentions all violations
        ("say hello for the win", False, " a"),  # Mentions all violations
        ("say hello for the win", False, "the win"),  # Mentions all violations
    ],
)
def test_check_banned_substring(text, expected, partial_error_message):
    restriction_fn = lambda x: restrictions.check_banned_substring(
        x, ["hello", "a", "the win"]
    )
    helper_test_restriction(None, text, expected, partial_error_message, restriction_fn)


ALPHA_SPACE_PUNCT_ENGLISH_ERR = restrictions.CHAR_SET_DICT[
    "alpha_space_punct_english"
].message
ENGLISH_ERR = "These are not valid English words"

S_ERR = "do not start with"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("snake slip slide", True, None),  # succeeds
        ("Snake Slip SLIDE", True, None),  # Okay with capital s
        ("snake volcano slide", False, S_ERR),  # all words must start with a/s/p/i/m
        ("slip snotaword slide", False, ENGLISH_ERR),  # errors for non-words
        ("snake, 'slip' slide.", True, None),  # handles various forms of punctuation
        ("something\nweird\nhere", False, S_ERR),  # No chaining with symbols
        ("something.weird.here", False, S_ERR),  # No chaining with symbols
        ("7", False, ALPHA_SPACE_PUNCT_ENGLISH_ERR),  # Numbers don't count
        (",,,,,,,,,,,,,........", True, None),  # Punctuation only okay
        ("snake...", True, None),  # Excess punctuation okay
        ("... --- ...", True, None),  # Punctuation only okay
    ],
)
def test_starts_with_aspmi(text, expected, partial_error_message):
    helper_test_restriction(
        "valid_words_starts_with_aspmi", text, expected, partial_error_message
    )


UNIQUE_ERR = " are not unique"
ALPHA_SPACE_ENGLISH_ERR = restrictions.CHAR_SET_DICT["alpha_space_english"].message


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("apple banana cherry", True, None),  # succeeds
        ("apple apple banana", False, UNIQUE_ERR),  # fails for repeat
        ("banana apple apple banana", False, "apple"),  # mentions both repeats
        ("banana apple apple banana", False, "banana"),  # mentions both repeats
        (
            "apple 'apple' banana",
            False,
            ALPHA_SPACE_ENGLISH_ERR,
        ),  # punctuation not allowed
        ("apple APPLE banana", False, UNIQUE_ERR),  # fails for repeats with diff caps
        ("appleapple", False, ENGLISH_ERR),  # fails for non-word
    ],
)
def test_valid_unique_words(text, expected, partial_error_message):
    helper_test_restriction("valid_unique_words", text, expected, partial_error_message)


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("pythn", True, None),  # No vowel
        ("python", False, BANNED_ERR),  # Has vowel
        ("PYTHON", False, BANNED_ERR),  # Has vowel (capitalized)
        ("pythan", False, BANNED_ERR),  # Has vowel
        ("pythin", False, BANNED_ERR),  # Has vowel
        ("pythen", False, BANNED_ERR),  # Has vowel
        ("pythun", False, BANNED_ERR),  # Has vowel
        ("pythön", True, None),  # Variants are OK (since diff tokenization)
        ("pythón", True, None),  # Variants on vowels are OK
        ("你好", True, None),  # diff alphabets okay
        ("𝕋𝕖𝕩𝕥𝔼𝕕𝕚𝕥𝕠𝕣", True, None),  # special characters okay
        ("?!,", True, None),  # punctuation okay
    ],
)
def test_no_vowels(text, expected, partial_error_message):
    helper_test_restriction("no_vowels", text, expected, partial_error_message)


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("pythn", True, None),  # No vowel
        ("python", False, BANNED_ERR),  # Has vowel
        ("PYTHON", False, BANNED_ERR),  # Has vowel (capitalized)
        ("pythan", False, BANNED_ERR),  # Has vowel
        ("pythin", False, BANNED_ERR),  # Has vowel
        ("pythen", False, BANNED_ERR),  # Has vowel
        ("pythun", False, BANNED_ERR),  # Has vowel
        ("pythön", False, ALNUM_SPACE_PUNCT_ERR),  # Variants are banned
        ("你好", False, ALNUM_SPACE_PUNCT_ERR),  # diff alphabest banned
        ("𝕋𝕖𝕩𝕥𝔼𝕕𝕚𝕥𝕠𝕣", False, ALNUM_SPACE_PUNCT_ERR),  # error on special characters
        ("?!,", True, None),  # punctuation ok
    ],
)
def test_alphanumeric_no_vowels(text, expected, partial_error_message):
    helper_test_restriction(
        "alphanumeric_no_vowels", text, expected, partial_error_message
    )


SPACE_ERR = "You cannot use whitespace characters"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("", True, None),  # Empty is fine
        ("a b", False, SPACE_ERR),  # Has space
        ("a\tb", False, SPACE_ERR),  # Has tab
        ("a\nb", False, SPACE_ERR),  # Has newline
        ("pythön", True, None),  # Variants okay
        ("你好", True, None),  # diff alphabets okay
        ("𝕋𝕖𝕩𝕥𝔼𝕕𝕚𝕥𝕠𝕣", True, None),  # special characters okay
        ("?!,", True, None),  # punctuation ok
    ],
)
def test_no_whitespace_any_characters(text, expected, partial_error_message):
    helper_test_restriction(
        "no_whitespace_any_characters", text, expected, partial_error_message
    )


ALNUM_PUNCT_ERR = restrictions.CHAR_SET_DICT["alnum_punct"].message


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("", True, None),  # Empty is fine
        ("a b", False, ALNUM_PUNCT_ERR),  # Has space
        ("a\tb", False, ALNUM_PUNCT_ERR),  # Has tab
        ("a\nb", False, ALNUM_PUNCT_ERR),  # Has newline
        ("pythön", False, ALNUM_PUNCT_ERR),  # Variants okay
        ("你好", False, ALNUM_PUNCT_ERR),  # diff alphabets banned
        ("𝕋𝕖𝕩𝕥𝔼𝕕𝕚𝕥𝕠𝕣", False, ALNUM_PUNCT_ERR),  # special characters banned
        ("?!,", True, None),  # punctuation ok
    ],
)
def test_no_whitespace_alphanum_punc(text, expected, partial_error_message):
    helper_test_restriction(
        "no_whitespace_alphanum_punc", text, expected, partial_error_message
    )


IS_ENGLISH_ERR = "is an English word"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("aqppl", True, None),  # Invalid words are okay
        ("apple", False, IS_ENGLISH_ERR),  # Valid words are not okay
        ("apple!", False, IS_ENGLISH_ERR),  # Valid words with punctuation are not okay
        ("'apple'", False, IS_ENGLISH_ERR),  # Valid words with quotes are not okay
        ("apple?!?", False, IS_ENGLISH_ERR),  # Valid words with weird punctuation bad
        (
            "apple-",
            False,
            IS_ENGLISH_ERR,
        ),  # Valid words with weird punctuation are not okay
        (
            ".apple",
            False,
            IS_ENGLISH_ERR,
        ),  # Valid words with weird punctuation are not okay
        (
            "~apple",
            False,
            IS_ENGLISH_ERR,
        ),  # Valid words with weird punctuation are not okay
        (
            "APPLE",
            False,
            IS_ENGLISH_ERR,
        ),  # Valid words with weird capitalization banned
        ("aPpLe", False, IS_ENGLISH_ERR),  # ignores capitalization
        ("p.e.t", False, IS_ENGLISH_ERR),  # Letters are words
        ("penguino", True, None),  # Other languages are OK
        ("... --- ...", True, None),  # Encodings are OK
        ("授予访问权限", True, None),  # Other alphabets OK
    ],
)
def test_no_english_words(text, expected, partial_error_message):
    helper_test_restriction("no_english_words", text, expected, partial_error_message)


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("aqppl", True, None),  # Invalid words are okay
        ("apple", False, IS_ENGLISH_ERR),  # Valid words are not okay
        ("apple!", False, IS_ENGLISH_ERR),  # Valid words with punctuation are not okay
        ("'apple'", False, IS_ENGLISH_ERR),  # Valid words with quotes are not okay
        ("apple?!?", False, IS_ENGLISH_ERR),  # Valid words with weird punctuation bad
        ("apple-", False, IS_ENGLISH_ERR),  # Valid words with weird punctuation bad
        (".apple", False, IS_ENGLISH_ERR),  # Valid words with weird punctuation bad
        ("~apple", False, IS_ENGLISH_ERR),  # Valid words with weird punctuation bad
        ("APPLE", False, IS_ENGLISH_ERR),  # Valid words with weird capitalization bad
        ("aPpLe", False, IS_ENGLISH_ERR),  # ignores capitalization
        ("p.e.t", False, IS_ENGLISH_ERR),  # Letters are words
        ("penguino", True, None),  # Other languages are OK
        ("... --- ...", True, None),  # Encodings are OK
        ("授予访问权限", False, ALNUM_SPACE_PUNCT_ERR),  # Other alphabets banned
    ],
)
def test_no_english_words_alphanumeric(text, expected, partial_error_message):
    helper_test_restriction(
        "no_english_words_alphanumeric", text, expected, partial_error_message
    )


LEN_ERR = "longer than 5 characters"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("apple sauce is yummy", True, None),  # Succeeds for short words
        ("applesauce is yummy", False, LEN_ERR),  # Fails for long words
        ("app les auce", False, ENGLISH_ERR),  # Can't cheat by adding spaces
        ("a p p l e s a u c e", True, None),  # I guess this is okay?
        ("asdf", False, ENGLISH_ERR),  # Fails for invalid words
        ("is!", False, ALPHA_SPACE_ENGLISH_ERR),  # Punct not allowed
    ],
)
def test_valid_words_length_5(text, expected, partial_error_message):
    helper_test_restriction(
        "valid_words_leq_length_5", text, expected, partial_error_message
    )


LEN_ERR = "shorter than 8 characters"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("shortest apartment", True, None),  # Succeeds for long words
        ("shorter apartment", False, LEN_ERR),  # Fails for short words
        (
            "daisy-chain",
            False,
            ALPHA_SPACE_ENGLISH_ERR,
        ),  # Words separated by punct fail
        ("asdfghjkl", False, ENGLISH_ERR),  # Fails for invalid words
    ],
)
def test_valid_words_length_8(text, expected, partial_error_message):
    helper_test_restriction(
        "valid_words_geq_length_8", text, expected, partial_error_message
    )


RELATED_ERR = "words related to"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("hi", True, None),  # unrelated words okay
        ("entry hi", False, RELATED_ERR),  # related words banned
        ("EnTrY", False, RELATED_ERR),  # weirdly formatted related words banned
        ("penguino", False, ENGLISH_ERR),  # non-english words banned
    ],
)
def test_valid_no_related_words(text, expected, partial_error_message):
    helper_test_restriction(
        "valid_no_related_words", text, expected, partial_error_message
    )


def test_extract_words():
    # Test we split on all allowed punctuation
    punctuation_list = '!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~'
    whitespace_list = " \t\n"
    for p in punctuation_list + whitespace_list:
        s = f"hello{p}world"
        split = ["hello", "world"]
        assert restrictions.extract_words(s) == split
    assert restrictions.extract_words("hello") == ["hello"]
    assert restrictions.extract_words("hello world") == [
        "hello",
        "world",
    ]
    assert restrictions.extract_words("hello    world") == [
        "hello",
        "world",
    ]
    assert restrictions.extract_words("hello, world") == [
        "hello",
        "world",
    ]
    assert restrictions.extract_words("hello,,?,world") == [
        "hello",
        "world",
    ]

    # Apostrophes only split when they're between letters
    assert restrictions.extract_words("hello'world") == ["hello'world"]
    assert restrictions.extract_words("hello' world") == [
        "hello",
        "world",
    ]


DISTINCT_ERR = "distinct words"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        ("one two three four five six five six", True, None),  # Succeeds 6 unique words
        (
            "one two three four five six seven",
            False,
            DISTINCT_ERR,
        ),  # Fails for 7 unique words
        (
            "one two three four five six FIVE Six",
            True,
            None,
        ),  # CAPS counts as same word
        (
            'one, two! "three" four five/six five six',
            True,
            None,
        ),  # Punctuation counts as same word
        (
            "one two's three four five six five six",
            True,
            None,
        ),  # Apostrophe doesn't split words
        (
            "one's two's three four five one two",
            False,
            DISTINCT_ERR,
        ),  # Apostrophe doesn't split words
        (
            "one-two.three,four$five*six?seven",
            False,
            DISTINCT_ERR,
        ),  # Other punctuation does split words
        (
            "onetwothreefourfivesixfivesix",
            False,
            ENGLISH_ERR,
        ),  # Must be valid English words
        ("pythón", False, ENGLISH_ERR),  # Must be valid English words
    ],
)
def test_unique_words_6(text, expected, partial_error_message):
    helper_test_restriction("unique_words_6", text, expected, partial_error_message)


DISTINCT_ERR = "distinct words"

THREE_LINES_ERR = "3 lines"
SYLLABLES_ERR = "syllables"


@pytest.mark.parametrize(
    "text, expected, partial_error_message",
    [
        (
            "i am a person \n you are a person as well \n weird coincidence",
            True,
            None,
        ),  # Valid Haiku
        (
            "i am a person? $$$$$$ \n you ARE a PERSON as Well \n weird 'coincidence' 🌕🌕🌕",
            True,
            None,
        ),  # Valid Haiku with weird caps and punc and emojis
        (
            "i am a person? $$$$$$ \n you ARE a PERSON as Well \n weird 'coincidence' 𝕋𝕖𝕩𝕥",
            False,
            NOT_ENGLISH_ERR,
        ),  # Valid Haiku with special characters that are treated as 'words'
        (" line one \n line two  ", False, THREE_LINES_ERR),  # Too few lines
        (
            " line one \n line two \n line three \n line four ",
            False,
            THREE_LINES_ERR,
        ),  # Too many lines
        (
            "i am a person \n you are a person as well \n weird coincidences",
            False,
            SYLLABLES_ERR,
        ),  # Wrong syllable count
        (
            "twas brillig, and the \n slithy toves did gyre \n and gimble in the",
            False,
            ENGLISH_ERR,
        ),  # Valid Haiku, but with non-dictionary words
    ],
)
def test_haiku(text, expected, partial_error_message):
    helper_test_restriction("haiku", text, expected, partial_error_message)


def test_long_restrictions():
    # Time any restrictions where computing them might take a while.
    max_time_cutoff = 0.005  # 5 ms
    num_repeats = 10

    def check_fast(fn, input_str):
        assert (
            timeit.timeit(lambda: fn(input_str), number=num_repeats) / num_repeats
            < max_time_cutoff
        )

    # Long words
    check_fast(
        restrictions.extract_words,
        " deoxyribonucleoprotein" * settings.MAX_ATTACK_LENGTH_TOKENS,
    )
    # One long word
    check_fast(restrictions.extract_words, "a" * settings.MAX_ATTACK_LENGTH_TOKENS * 5)
    # Weird tokens
    check_fast(restrictions.extract_words, "🌕" * settings.MAX_ATTACK_LENGTH_TOKENS)
    # All punctuation
    check_fast(restrictions.extract_words, "." * settings.MAX_ATTACK_LENGTH_TOKENS * 5)

    def check_restriction_fast(restriction, input_str):
        assert (
            timeit.timeit(lambda: restriction.is_valid(input_str), number=num_repeats)
            / num_repeats
            < max_time_cutoff
        )

    # HAIKU (includes syllable counting)
    haiku_restriction = RESTRICTIONS_DICT["haiku"]

    # We'll make all of these 3 lines long so it doesn't get rejected early for not having enough lines
    # Try a long string made of "normal" words, in the cmudict
    long_string = "a\na\n" + (" hello" * settings.MAX_ATTACK_LENGTH_TOKENS)
    check_restriction_fast(haiku_restriction, long_string)

    # Try a string made of long words which are not in the dictionary
    long_string = "a\na\n" + (
        " deoxyribonucleoprotein" * settings.MAX_ATTACK_LENGTH_TOKENS
    )
    check_restriction_fast(haiku_restriction, long_string)

    # Non-words should get rejected but let's check them just in case
    long_string = "a\na\n" + (
        " asdfkarjlfasdadfwefqwe" * settings.MAX_ATTACK_LENGTH_TOKENS
    )
    check_restriction_fast(haiku_restriction, long_string)
    long_string_one_word = "a\na\n" + (
        "asdfkarjlfasdadfwefqwe" * settings.MAX_ATTACK_LENGTH_TOKENS
    )
    check_restriction_fast(haiku_restriction, long_string_one_word)
    long_string_emojis = "a\na\n" + ("🌕" * settings.MAX_ATTACK_LENGTH_TOKENS)
    check_restriction_fast(haiku_restriction, long_string_emojis)


def test_char_sets():
    # Mapping from "keyword" (in the char set message) to a character which should be valid iff the keyword is present
    char_set_dict = {
        "a-z": "a",
        "numbers": "7",
        "whitespace": " ",
        "punctuation": ".",
        "symbols": "🌕",
    }
    other_symbols = "你ö𝕋"

    for char_set_name, char_set in restrictions.CHAR_SET_DICT.items():
        for keyword, allowed_char in char_set_dict.items():
            if keyword in char_set.message:
                assert char_set.check_fn(
                    allowed_char
                ).is_valid, f"{char_set_name} doesn't allow {allowed_char}"
            else:
                assert not char_set.check_fn(
                    allowed_char
                ).is_valid, f"{char_set_name} allows {allowed_char}"
        # all input restrictions ban the other symbols
        for disallowed_char in other_symbols:
            assert not char_set.check_fn(
                disallowed_char
            ).is_valid, f"{char_set_name} allows {disallowed_char}"
