import itertools
import random
import re
import unicodedata
from collections import defaultdict
from dataclasses import dataclass
from enum import Enum
from functools import partial
from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Union

import cmudict
import syllables

# This is a hack to avoid circular imports
# while still having type checking for Account
if TYPE_CHECKING:
    from promptgame.gameui.models import Account

from django.apps.registry import Apps
from django.conf import settings

from promptgame.gameui.tiers import tier_index
from promptgame.gameui.utils import get_number_tokens

RESTRICTIONS_DICT = {}

# TODO(olivia): consider using alt format where we search for disallowed chars
ALNUM_SPACE_MATCHER = re.compile(r"^[a-zA-Z0-9\s\n\t]*$")
ALPHA_SPACE_MATCHER = re.compile(r"^[a-zA-Z\s\n\t]*$")
ALNUM_SPACE_PUNCT_MATCHER = re.compile(
    r"^[a-zA-Z0-9\s\n\t!\"#$%&\'()*+,-./:;<=>?@\[\]^_`{|}~]*$"
)
ALPHA_SPACE_PUNCT_MATCHER = re.compile(
    r"^[a-zA-Z\s\n\t!\"#$%&\'()*+,-./:;<=>?@\[\]^_`{|}~]*$"
)
SPACE_MATCHER = re.compile(r"[\s\u200B\u2009]")
ALNUM_PUNCT_MATCHER = re.compile(r'^[a-zA-Z0-9!\\"#$%&\'()*+,-./:;<=>?@[\]^_`{|}~]*$')
# Match all runs of word chars (except underscore) and apostrophes
# (works in every script, not just English)
ALPHANUMERIC_AND_APOSTROPHE_RE = re.compile(r"(?P<word>([^\W_]|')+)")

EASY_WEIGHT = 10
MEDIUM_WEIGHT = 5
HARD_WEIGHT = 1
# Approximate Target: 70% easy, 25% medium, 5% hard


def describe_char(char):
    """Describe a character in a way that's useful for debugging."""
    desc = f"'{char}'"
    if (name := unicodedata.name(char, None)) is not None:
        desc += f" ({name.title()})"
    return desc


def zip_safe(*args):
    """Like zip, but first checks the zipped args are all the same length."""
    args_list = list(args)
    arg_lengths = [len(arg) for arg in args_list]
    if not all([arg_lengths[0] == arg_len for arg_len in arg_lengths]):
        raise ValueError(f"Not all args are the same length: {arg_lengths}")
    return zip(*args)


def join_iters(iters: Iterable[Iterable[str]]) -> Iterable[str]:
    """Takes an iterable of iterables of strings, and joins together the strings in each sub-iterable.
    Useful when working with itertools + strings.
    Example: join_iters([['a', 'b'], ['c', 'd']]) -> Generator['ab', 'cd']

    """
    for str_iter in iters:
        yield "".join(str_iter)


def string_combinations(s: str, r: int) -> Iterable[str]:
    """All ways of choosing r distinct characters from s (will not double-count for different orderings;
    e.g. string_combinations("ab", 2) only gives "ab", and not "ba").
    """
    return join_iters(itertools.combinations(s, r))


def string_product(*iters: Iterable[str]) -> Iterable[str]:
    """ "All combinations of one character from the first iterable, one from the second, one from the third, etc."""
    return join_iters(itertools.product(*iters))


@dataclass
class RestrictionResult:
    is_valid: bool
    error_message: Optional[str] = None


class Difficulty(Enum):
    EASY = 0
    MEDIUM = 1
    HARD = 2
    DEFAULT = MEDIUM


def get_restriction_difficulty(account: "Account") -> Difficulty:
    restriction = RESTRICTIONS_DICT[account.restriction_id]
    return restriction.difficulty


def get_higher_difficulty(account: "Account") -> Difficulty:
    restriction = RESTRICTIONS_DICT[account.restriction_id]
    if restriction.difficulty == Difficulty.HARD:  # No higher difficulty bucket
        return Difficulty.HARD
    return Difficulty(restriction.difficulty.value + 1)


def get_lower_difficulty(account: "Account") -> Difficulty:
    restriction = RESTRICTIONS_DICT[account.restriction_id]
    if restriction.difficulty == Difficulty.EASY:  # No lower difficulty bucket
        return Difficulty.EASY
    return Difficulty(restriction.difficulty.value - 1)


class Restriction:
    def __init__(
        self,
        id: str,
        description: str,
        function: Callable[[str], RestrictionResult],
        difficulty: Difficulty,
        restriction_group: Optional[str] = None,
    ):
        self.id = id
        self.description = description
        self.function = function
        self.difficulty = difficulty
        if restriction_group is None:
            restriction_group = id
        self.restriction_group = restriction_group

    def is_valid(self, user_text: str) -> RestrictionResult:
        return self.function(user_text)


def get_restriction(account: "Account") -> Restriction:
    return RESTRICTIONS_DICT[account.restriction_id]


def check_restriction_enabled(
    attacking_account: "Account", defending_account: "Account"
):
    # Restrictions apply if an account attacks an account in a lower tier.
    return tier_index(attacking_account.balance) > tier_index(defending_account.balance)


# Below, we register all the different restrictions.


def run_on_import(fn):
    """This module is organized into top-level functions decorated with `@run_on_import`.
    This helps us run lots of code to register restrictions without polluting the global
    namespace."""
    fn()
    return fn


def check_token_limit(user_text: str, token_limit: int) -> RestrictionResult:
    token_result = get_number_tokens(user_text, token_limit)
    if not token_result.verification_passed:
        return RestrictionResult(
            False,
            f"Your attack used more than {token_limit} tokens.",
        )
    return RestrictionResult(True)


def is_valid_english_word(word: str) -> bool:
    return word.lower() in ENGLISH_WORDS


def split_words_on_whitespace(s: str) -> List[str]:
    s = s.strip().lower()
    # Split on whitespace, drop empty strings
    words = [w for w in s.split() if len(w) > 0]
    return words


def extract_words(s: str) -> List[str]:
    words = (
        m.groupdict()["word"]
        for m in ALPHANUMERIC_AND_APOSTROPHE_RE.finditer(s.lower())
    )
    # Since aposrophes don't split words, trim off starting and ending
    # apostrophes
    stripped = (w.strip("'") for w in words)
    return [w for w in stripped if len(w) > 0]


def get_word_violations(s: str, valid_fn: Callable[[str], bool]) -> List[str]:
    words = extract_words(s)
    bad_words = {}
    for word in words:
        if not valid_fn(word):
            bad_words[word] = True
    return list(bad_words.keys())


def check_valid_english_words(s: str) -> RestrictionResult:
    bad_words = get_word_violations(s, is_valid_english_word)
    if len(bad_words) == 0:
        return RestrictionResult(True)
    return RestrictionResult(
        False, f"These are not valid English words: {list_to_str(bad_words)}."
    )


def check_char_set_match(
    user_text: str, char_set_matcher: re.Pattern, error_message: str
) -> RestrictionResult:
    for char in user_text:
        if not char_set_matcher.fullmatch(char):
            return RestrictionResult(
                False, f"{error_message} Violating characters: {describe_char(char)}."
            )
    return RestrictionResult(True)


@dataclass
class CharSet:
    message: str
    check_fn: Callable[[str], RestrictionResult]


CHAR_SET_DICT = {}


@run_on_import
def register_char_sets() -> None:
    def register_input_restriction(
        char_set_name: str,
        input_check_fn: Callable[[str], RestrictionResult],
        message: str,
    ) -> None:
        """Put a new input_check_fn in the CHAR_SET_DICT."""
        CHAR_SET_DICT[char_set_name] = CharSet(
            message=message,
            check_fn=input_check_fn,
        )

    def register_char_set_with_pattern(
        char_set_name: str, char_set_matcher: re.Pattern, message: str
    ) -> None:
        """Turn a char_set_matcher into an input check fn and store it."""
        fn = partial(
            check_char_set_match,
            char_set_matcher=char_set_matcher,
            error_message=message,
        )
        register_input_restriction(char_set_name, fn, message)

    def compose_char_check_english_check(
        char_set_check_fn: Callable[[str], RestrictionResult]
    ) -> Callable[[str], RestrictionResult]:
        """First check the input space is valid.
        Then check if the word is valid English if it passes."""

        def composed_fn(user_text: str) -> RestrictionResult:
            # First call the char set check
            result = char_set_check_fn(user_text)
            if not result.is_valid:
                return result
            return check_valid_english_words(user_text)

        return composed_fn

    def register_char_set_valid_english(
        char_set_name: str, char_set_matcher: re.Pattern, message: str
    ) -> None:
        """Turn a char_set_matcher into an input check fn which first
        checks the input space, then if it succeeds checks for English."""
        char_set_check_fn = partial(
            check_char_set_match,
            char_set_matcher=char_set_matcher,
            error_message=message,
        )
        composed_fn = compose_char_check_english_check(char_set_check_fn)
        register_input_restriction(char_set_name, composed_fn, message)

    # Register all the input spaces

    msg = "Use a-z, numbers, whitespace, and punctuation only."
    register_char_set_with_pattern("alnum_space_punct", ALNUM_SPACE_PUNCT_MATCHER, msg)

    msg = "Use a-z and whitespace only."
    register_char_set_with_pattern("alpha_space", ALPHA_SPACE_MATCHER, msg)

    msg = "Use a-z, whitespace, and punctuation only."
    register_char_set_with_pattern("alpha_space_punct", ALPHA_SPACE_PUNCT_MATCHER, msg)

    msg = "Use a-z, numbers, and punctuation only."
    register_char_set_with_pattern("alnum_punct", ALNUM_PUNCT_MATCHER, msg)

    msg = "Use valid English words containing a-z, whitespace, and punctuation only."
    register_char_set_valid_english(
        "alpha_space_punct_english", ALPHA_SPACE_PUNCT_MATCHER, msg
    )

    msg = "Use valid English words containing a-z and whitespace only."
    register_char_set_valid_english("alpha_space_english", ALPHA_SPACE_MATCHER, msg)

    msg = "Use valid English words containing a-z, whitespace, and symbols/punctuation only."
    # For this one, we actually only apply the english check.
    # The char set is defined by our English regex and word list.
    # It's not 100% accurate since our English dict includes some numbers, but it's close enough.
    register_input_restriction("english", check_valid_english_words, msg)


def check_banned_substring(
    user_text: str, banned_words: List[str]
) -> RestrictionResult:
    user_text = user_text.lower()
    bad_words = []
    for word in banned_words:
        if word in user_text:
            bad_words.append(word)
    if len(bad_words) > 0:
        return RestrictionResult(
            False, f"Using banned letters/words: {list_to_str(bad_words)}."
        )
    return RestrictionResult(True)


def char_list_to_str(words: Union[str, Iterable[str]]) -> str:
    """Takes in an iterable of strs (e.g. ['a', 'b', 'c']) and returns a string like "'a', 'b', or 'c'" """
    if len(words) == 1:
        return f"'{words[0]}'"
    elif len(words) == 2:
        return f"'{words[0]}' or '{words[1]}'"
    else:
        all_but_last = "', '".join(words[:-1])
        return f"'{all_but_last}' or '{words[-1]}'"


def list_to_str(words: Iterable[str]) -> str:
    """Takes in an iterable of strs (e.g. ['a', 'b', 'c']) and returns a string like " 'a', 'b', or 'c'" """
    return ", ".join(words)


ENGLISH_WORDS = set()


@run_on_import
def load_english_words() -> None:
    # Word list found at https://github.com/dwyl/english-words/blob/master/words_alpha.txt
    with open(settings.ENGLISH_WORDS_PATH) as f:
        for line in f:
            ENGLISH_WORDS.add(line.strip().lower())


@run_on_import
def register_tokens_restriction() -> None:
    # Max tokens restriction
    token_limits = [10, 20, 30, 40, 50]
    difficulties = [
        Difficulty.HARD,
        Difficulty.HARD,
        Difficulty.MEDIUM,
        Difficulty.EASY,
        Difficulty.EASY,
    ]
    for limit, difficulty in zip_safe(token_limits, difficulties):
        id = f"token_limit_{limit}"
        description = f"Your attack must be at most {limit} tokens."
        function = partial(check_token_limit, token_limit=limit)
        RESTRICTIONS_DICT[id] = Restriction(
            id, description, function, difficulty, restriction_group="token_limit"
        )


@run_on_import
def register_common_1k() -> None:
    id = "common_1k_words"
    char_set_common = CHAR_SET_DICT["alpha_space"]
    description = (
        "You can only use the 1,000 most common English words (word list from tinyurl.com/words1k). "
        + char_set_common.message
    )
    # List taken from https://gist.github.com/deekayen/4148741
    with open(settings.COMMON_WORDS_PATH) as f:
        common_words = f.read().splitlines()
    # lower case (e.g. I -> i)
    common_words = {word.lower() for word in common_words}

    def check_most_common_1k_words(user_text):
        # Throw an error if there are characters besides a-z and whitespace (including newlines)
        user_text = user_text.lower()
        result = char_set_common.check_fn(user_text)
        if not result.is_valid:
            return result

        bad_words = get_word_violations(user_text, lambda word: word in common_words)
        if len(bad_words) > 0:
            return RestrictionResult(
                False,
                f"These words are not in the 1,000 most common: {list_to_str(bad_words)}.",
            )
        return RestrictionResult(True)

    RESTRICTIONS_DICT[id] = Restriction(
        id, description, check_most_common_1k_words, difficulty=Difficulty.MEDIUM
    )

    # Also create one which does the reverse -- bans the top 1k words
    uncommon_id = "not_common_1k_words"
    char_set_uncommon = CHAR_SET_DICT["alnum_space_punct"]
    uncommon_description = (
        "You cannot use the 1,000 most common English words (word list from tinyurl.com/words1k). "
        + char_set_uncommon.message
    )

    def check_not_common_1k_words(user_text):
        # Throw an error if there are characters besides a-z and whitespace (including newlines)
        user_text = user_text.lower()
        result = char_set_uncommon.check_fn(user_text)
        if not result.is_valid:
            return result

        bad_words = get_word_violations(
            user_text, lambda word: word not in common_words
        )
        if len(bad_words) > 0:
            return RestrictionResult(
                False,
                f"These words are in the 1,000 most common: {list_to_str(bad_words)}.",
            )
        return RestrictionResult(True)

    RESTRICTIONS_DICT[uncommon_id] = Restriction(
        uncommon_id,
        uncommon_description,
        check_not_common_1k_words,
        difficulty=Difficulty.MEDIUM,
    )


@run_on_import
def register_no_vowels() -> None:
    no_vowels_id = "no_vowels"
    no_vowels_description = "You cannot use vowels (a, e, i, o, or u)."

    def check_no_vowels(user_text):
        return check_banned_substring(user_text, ["a", "e", "i", "o", "u"])

    RESTRICTIONS_DICT[no_vowels_id] = Restriction(
        no_vowels_id,
        no_vowels_description,
        check_no_vowels,
        difficulty=Difficulty.MEDIUM,
        restriction_group="no_vowels",
    )

    alnum_no_vowels_id = "alphanumeric_no_vowels"
    char_set = CHAR_SET_DICT["alnum_space_punct"]
    alnum_no_vowels_description = (
        "You cannot use vowels (a, e, i, o, or u). " + char_set.message
    )

    def alphanumeric_check_no_vowels(user_text):
        alphanumeric_result = char_set.check_fn(user_text)
        if not alphanumeric_result.is_valid:
            return alphanumeric_result
        return check_no_vowels(user_text)

    RESTRICTIONS_DICT[alnum_no_vowels_id] = Restriction(
        alnum_no_vowels_id,
        alnum_no_vowels_description,
        alphanumeric_check_no_vowels,
        difficulty=Difficulty.HARD,
        restriction_group="no_vowels",
    )


@run_on_import
def register_valid_starts_with_letter() -> None:
    def check_word_startswith(word, starting_letters):
        return any([word.startswith(letter) for letter in starting_letters])

    char_set = CHAR_SET_DICT["alpha_space_punct_english"]

    def check_starts_with_letters(user_text, starting_letters):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        bad_words = get_word_violations(
            user_text, lambda word: check_word_startswith(word, starting_letters)
        )
        if len(bad_words) > 0:
            letter_str = char_list_to_str(starting_letters)
            return RestrictionResult(
                False,
                f"These words do not start with {letter_str}: {list_to_str(bad_words)}.",
            )
        return RestrictionResult(True)

    # most common starting letters: spc (in dictionaries); atio (in normal text)
    # https://en.wikipedia.org/wiki/Letter_frequency

    # access granted starting letters
    access_granted_letters = "ag"

    # Easier starting letters
    easy_letter_list = "spct"

    # Harder starting letters (excluding ones with < 3% starting frequency in normal text
    # or < 3% in dictionaries)
    harder_letter_list = "bdfhmi"

    all_combos = list(
        string_product(
            string_combinations(
                access_granted_letters, 1
            ),  # iterable of all access granted letters
            string_combinations(
                easy_letter_list, 2
            ),  # iterable of all distinct combinations of 2 easy letters
            string_combinations(
                harder_letter_list, 2
            ),  # iterable of all distinct combinations of 2 hard letters
        )
    )  # All combos of 1 access granted letter, 2 easy letters, and 2 hard letters

    for letters in all_combos:
        id = f"valid_words_starts_with_{letters}"
        letter_str = char_list_to_str(letters)
        description = (
            f"You can only use words which start with the letters '{letter_str}'. "
            + char_set.message
        )
        RESTRICTIONS_DICT[id] = Restriction(
            id,
            description,
            partial(check_starts_with_letters, starting_letters=letters),
            difficulty=Difficulty.HARD,
            restriction_group="valid_words_starts_with",
        )


@run_on_import
def register_valid_contains_or_doesnt_contain_letters() -> None:
    char_set = CHAR_SET_DICT["alpha_space_punct_english"]

    def check_doesnt_contain(user_text, letter_list):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        error_messages = []
        for letter in letter_list:
            bad_words = get_word_violations(
                user_text, lambda word: not letter in word.lower()
            )
            if len(bad_words) > 0:
                error_messages.append(
                    f"These words contain '{letter}': {list_to_str(bad_words)}."
                )
        if len(error_messages) > 0:
            return RestrictionResult(
                False, f"Restriction violation: {' '.join(error_messages)}"
            )
        return RestrictionResult(True)

    def check_contains(user_text, letter_list):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        bad_words = get_word_violations(
            user_text, lambda word: any([letter in word for letter in letter_list])
        )
        if len(bad_words) > 0:
            letter_str = char_list_to_str(letter_list)
            return RestrictionResult(
                False,
                f"These words to not contain {letter_str}: {list_to_str(bad_words)}",
            )
        return RestrictionResult(True)

    def add_restrictions_for_letter_lists(
        letter_lists: List[str],
    ) -> None:
        """Add restrictions (all words must contain/all words must not contain a
        certain set of letters) for each set of letters in `letter_lists`."""
        for r_letters in letter_lists:
            # Doesn't include letters
            id = f"doesnt_include_{r_letters}"
            formatted_letters = char_list_to_str(r_letters)
            description = (
                f"You can only use words which don't include {formatted_letters}. "
                + char_set.message
            )
            RESTRICTIONS_DICT[id] = Restriction(
                id,
                description,
                partial(check_doesnt_contain, letter_list=r_letters),
                difficulty=Difficulty.EASY,
                restriction_group="doesnt_include",
            )

            # Does include letters
            id = f"includes_{r_letters}"
            letter_str = char_list_to_str(r_letters)
            description = (
                f"You can only use words which contain {letter_str}. "
                + char_set.message
            )
            RESTRICTIONS_DICT[id] = Restriction(
                id,
                description,
                partial(check_contains, letter_list=r_letters),
                difficulty=Difficulty.HARD,
                restriction_group="includes",
            )

    # Single letter restrictions for common letters and letters in "access granted"
    more_common_letters = "acesgrntdihlo"
    # Includes all letters which are not in the common letters and are not unusually rare (jkqxz)
    less_common_letters = "bfmpuvwy"
    # Check that the letters are disjoint
    assert len(set(more_common_letters).intersection(set(less_common_letters))) == 0

    # Single-letter restrictions
    add_restrictions_for_letter_lists(
        list(more_common_letters),
    )

    # Distinct letter pair restrictions
    letter_pairs = list(
        string_product(iter(more_common_letters), iter(less_common_letters))
    )
    add_restrictions_for_letter_lists(
        letter_pairs,
    )


@run_on_import
def register_valid_unique_words() -> None:
    id = "valid_unique_words"
    char_set = CHAR_SET_DICT["alpha_space_english"]
    description = "You can't use the same word twice. " + char_set.message

    def check_unique_words(user_text):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        words = extract_words(user_text)
        bad_words = {}
        past_words = set()
        for word in words:
            if word in past_words:
                bad_words[word] = True
            else:
                past_words.add(word)
        bad_words = list(bad_words.keys())
        if len(bad_words) > 0:
            return RestrictionResult(
                False, f"These words are not unique: {list_to_str(bad_words)}."
            )
        return RestrictionResult(True)

    RESTRICTIONS_DICT[id] = Restriction(
        id, description, check_unique_words, difficulty=Difficulty.MEDIUM
    )


@run_on_import
def register_no_whitespace() -> None:
    id = "no_whitespace_any_characters"
    description = "You cannot use whitespace characters, including newlines."

    def check_no_whitespace(user_text):
        # Use re to see if the input contains spaces, tabs, or newlines
        # Should match the following kinds of whitespace:
        # (space)
        # \t (tab)
        # \n (newline)
        # \r (carriage return)
        # \u200B (zero-width space)
        # \u2009 (thin space)
        # \f (form feed)
        # \v (vertical tab)
        if SPACE_MATCHER.search(user_text):
            return RestrictionResult(False, f"You cannot use whitespace characters.")
        return RestrictionResult(True)

    RESTRICTIONS_DICT[id] = Restriction(
        id,
        description,
        check_no_whitespace,
        difficulty=Difficulty.EASY,
        restriction_group="no_whitespace",
    )

    id = "no_whitespace_alphanum_punc"
    char_set = CHAR_SET_DICT["alnum_punct"]
    description = (
        "You cannot use whitespace characters, including newlines. " + char_set.message
    )

    RESTRICTIONS_DICT[id] = Restriction(
        id,
        description,
        char_set.check_fn,
        difficulty=Difficulty.EASY,
        restriction_group="no_whitespace",
    )


@run_on_import
def register_no_related_words() -> None:
    # TODO(olivia): Maybe these are too expansive. Maybe we should just do all the word variants o
    # the banned words themselves, and only use synonyms of the same part of speech for the others?
    access_synonyms = {
        "access",
        "accessed",
        "accesses",
        "accessing",
        "accessor",
        "accessors",
        "enter",
        "entered",
        "enters",
        "entering",
        "entry",
        "entries",
        "entrance",
        "entrances",
        "approach",
        "approached",
        "approaches",
        "approaching",
        "admission",
        "admissions",
        "admit",
        "admits",
        "admitted",
        "admitting",
        "admittance",
        "admittances",
        "ingress",
        "ingresses",
        "ingression",
        "ingressions",
        "ingressing",
        "ingressed",
    }
    granted_synonyms = {
        "granted",
        "grant",
        "grants",
        "granting",
        "allow",
        "allows",
        "allowed",
        "allowing",
        "permit",
        "permits",
        "permitted",
        "permitting",
        "authorize",
        "authorizes",
        "authorized",
        "authorizing",
        "give",
        "gives",
        "gave",
        "given",
        "giving",
    }
    password_synonyms = {
        "password",
        "passwords",
        "pass word",
        "pass words",
        "pass-word",
        "pass-words",
        "code",
        "codes",
        "key",
        "keys",
        "combination",
        "combinations",
        "identification",
        "identifications",
        "id",
        "ids",
    }
    instruction_synonyms = {
        "instruction",
        "instructions",
        "instruct",
        "instructs",
        "instructed",
        "instructing",
        "direction",
        "directions",
        "direct",
        "directs",
        "directed",
        "directing",
        "order",
        "orders",
        "ordered",
        "ordering",
        "command",
        "commands",
        "commanded",
        "commanding",
        "rule",
        "rules",
        "ruled",
        "ruling",
        "guidance",
        "guidances",
        "guide",
        "guides",
        "guided",
        "guiding",
        "advice",
        "advices",
        "advise",
        "advises",
        "advised",
        "advising",
        "suggestion",
        "suggestions",
        "suggest",
        "suggests",
        "suggested",
        "suggesting",
        "mandate",
        "mandates",
        "mandated",
        "mandating",
        "goal",
        "goals",
        "goal",
        "goals",
    }
    door_synonyms = {
        "door",
        "doors",
        "gate",
        "gates",
        "portal",
        "portals",
        "entryway",
        "entryways",
        "entrance",
        "entrances",
    }

    def helper_check_related_words(words, banned_word, related_words):
        bad_words = get_word_violations(words, lambda word: word not in related_words)
        if len(bad_words) > 0:
            return RestrictionResult(
                False,
                f"Restriction violation: words related to {banned_word}: {list_to_str(bad_words)}.",
            )
        return RestrictionResult(True)

    char_set = CHAR_SET_DICT["alpha_space_punct_english"]
    id = "valid_no_related_words"
    description = (
        "You can only use words which are not synonyms or near-synonyms of 'access', 'granted', 'password', 'instruction', or 'door'. "
        + char_set.message
    )

    def check_no_related_words(user_text):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        banned_word_and_list = [
            ("access", access_synonyms),
            ("granted", granted_synonyms),
            ("password", password_synonyms),
            ("instruction", instruction_synonyms),
            ("door", door_synonyms),
        ]
        error_messages = []
        for banned_word, related_words in banned_word_and_list:
            resp = helper_check_related_words(user_text, banned_word, related_words)
            if not resp.is_valid:
                error_messages.append(resp.error_message)
        if len(error_messages) > 0:
            error_str = " ".join(error_messages)
            return RestrictionResult(False, f"Restriction violation: {error_str}")
        return RestrictionResult(True)

    RESTRICTIONS_DICT[id] = Restriction(
        id, description, check_no_related_words, difficulty=Difficulty.EASY
    )


@run_on_import
def register_word_length() -> None:
    char_set = CHAR_SET_DICT["alpha_space_english"]

    def check_words_length(user_text, length, leq):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result

        if leq:
            valid_fn = lambda word: len(word) <= length
            violation_word = "longer"
        else:
            valid_fn = lambda word: len(word) >= length
            violation_word = "shorter"

        bad_words = get_word_violations(user_text, valid_fn)
        if len(bad_words) > 0:
            return RestrictionResult(
                False,
                f"Restriction violation: these words are {violation_word} than {length} characters: {list_to_str(bad_words)}.",
            )
        return RestrictionResult(True)

    leq_lengths = [4, 5, 6]
    difficulties = [Difficulty.HARD, Difficulty.HARD, Difficulty.MEDIUM]
    for length, difficulty in zip_safe(leq_lengths, difficulties):
        id = f"valid_words_leq_length_{length}"
        description = (
            f"You can only use words which are {length} characters or less. "
            + char_set.message
        )
        RESTRICTIONS_DICT[id] = Restriction(
            id,
            description,
            partial(check_words_length, length=length, leq=True),
            difficulty=difficulty,
            restriction_group="valid_words_leq_length",
        )

    geq_lengths = [5, 6, 7, 8]
    for length in geq_lengths:
        id = f"valid_words_geq_length_{length}"
        description = (
            f"You can only use valid English words which are {length} characters or more. "
            + char_set.message
        )
        RESTRICTIONS_DICT[id] = Restriction(
            id,
            description,
            partial(check_words_length, length=length, leq=False),
            difficulty=Difficulty.MEDIUM,
            restriction_group="valid_words_geq_length",
        )


@run_on_import
def register_no_english_words() -> None:
    # No English words
    id = "no_english_words"
    description = "You cannot use any valid English words."

    def check_no_english_words(user_text):
        words = extract_words(user_text)
        for word in words:
            if is_valid_english_word(word):
                return RestrictionResult(False, f"{word} is an English word.")
        return RestrictionResult(True)

    RESTRICTIONS_DICT[id] = Restriction(
        id,
        description,
        check_no_english_words,
        difficulty=Difficulty.EASY,
        restriction_group="no_english_words",
    )

    # No English words
    id = "no_english_words_alphanumeric"
    char_set = CHAR_SET_DICT["alnum_space_punct"]
    description = "You cannot use any valid English words. " + char_set.message

    def check_no_english_words_alphanumeric(user_text):
        res = char_set.check_fn(user_text)
        if not res.is_valid:
            return res
        return check_no_english_words(user_text)

    RESTRICTIONS_DICT[id] = Restriction(
        id,
        description,
        check_no_english_words_alphanumeric,
        difficulty=Difficulty.EASY,
        restriction_group="no_english_words",
    )


@run_on_import
def register_count_unique_words() -> None:
    char_set = CHAR_SET_DICT["english"]

    def check_unique(user_text, max_count):
        valid_result = char_set.check_fn(user_text)
        if not valid_result.is_valid:
            return valid_result
        num_words_used = len(set(extract_words(user_text)))
        if num_words_used > max_count:
            return RestrictionResult(
                False,
                f"Restriction violation: you used {num_words_used} distinct words, must use at most {max_count},",
            )
        return RestrictionResult(True)

    lengths = [5, 6, 7, 8]
    for length in lengths:
        id = f"unique_words_{length}"
        description = (
            f"You can use at most {length} distinct words. " + char_set.message
        )
        RESTRICTIONS_DICT[id] = Restriction(
            id,
            description,
            partial(check_unique, max_count=length),
            difficulty=Difficulty.HARD,
            restriction_group="unique_words",
        )


@run_on_import
def register_haiku() -> None:
    id = "haiku"
    char_set = CHAR_SET_DICT["english"]
    description = (
        "The attack must be a valid haiku (3 lines, syllable count 5/7/5). "
        + char_set.message
    )

    # import time
    # start = time.time()
    syllables_dict = cmudict.dict()
    # load_dict_time = time.time() - start
    # print(f'Loading cmudict took {load_dict_time} seconds.')

    def count_syllables(word):
        # We have two ways of counting syllables. cmudict has a dictionary which is
        # more accurate for words in the dictionary, but fails completely for words not
        # in the dictionary. syllables is less accurate (I think), but tries its best on
        # weird words.
        if word in syllables_dict:
            # https://stackoverflow.com/questions/405161/detecting-syllables-in-a-word/4103234#4103234
            # cmudict returns a list of the phonemes in the word. Vowels contain a stress marker indicated
            # by a digit on the end. We can extract these digits to count syllables.
            # If there are multiple pronounciations, we just take the first for simplicity.
            return len([y for y in syllables_dict[word.lower()][0] if y[-1].isdigit()])
        return syllables.estimate(word)

    def check_haiku(user_text):
        res_result = char_set.check_fn(user_text)
        if not res_result.is_valid:
            return res_result

        lines = [l for l in user_text.split("\n") if len(l) > 0]
        if not len(lines) == 3:
            return RestrictionResult(False, "A haiku must be 3 lines long.")
        line_counts = [5, 7, 5]
        for i in range(3):
            num_syllables = 0
            for word in extract_words(lines[i]):
                num_syllables += count_syllables(word)
                # Break early if we've already exceeded the syllable count
                # so it isn't slower for long strings
                if num_syllables > line_counts[i]:
                    break
            if not num_syllables == line_counts[i]:
                return RestrictionResult(
                    False,
                    f"Line {i+1} is {num_syllables} syllables but should be {line_counts[i]} syllables long.",
                )
        return RestrictionResult(True, None)

    RESTRICTIONS_DICT[id] = Restriction(
        id, description, check_haiku, difficulty=Difficulty.MEDIUM
    )


# mapping from Difficulty to a set of restriction_groups which contain
# at least one restriction of that difficulty.
DIFFICULTY_GROUP_MAP = defaultdict(lambda: defaultdict(lambda: []))
for restriction_id, restriction in RESTRICTIONS_DICT.items():
    DIFFICULTY_GROUP_MAP[restriction.difficulty][restriction.restriction_group].append(
        restriction
    )


def get_random_restriction_id(difficulty: Difficulty = Difficulty.DEFAULT) -> str:
    restriction = random.choice(  # choose a restriction within the group
        random.choice(  # choose a group of restrictions at this difficulty level
            list(DIFFICULTY_GROUP_MAP[difficulty].values())
        )
    )
    return restriction.id


def migration_make_all_account_restrictions_valid(apps: Apps, _) -> None:
    Account = apps.get_model("gameui", "Account")
    for account in Account.objects.iterator(chunk_size=128):
        if not account.restriction_id in RESTRICTIONS_DICT:
            account.restriction_id = get_random_restriction_id()
            account.save()


def migration_refresh_all_account_restrictions(apps: Apps, _) -> None:
    Account = apps.get_model("gameui", "Account")
    for account in Account.objects.iterator(chunk_size=128):
        # Start with a mix of EASY and MEDIUM restrictions
        difficulty = random.choice([Difficulty.EASY, Difficulty.MEDIUM])
        account.restriction_id = get_random_restriction_id(difficulty=difficulty)
        account.save()
