import sys
import string
from trie import TypeaheadSearchTrie


class Entry(object):
    """Simple container class for data entries."""
    def __init__(self, type, id, score, data):
        self.type = type
        self.id = id
        self.score = score
        self.data = data


class TypeAheadSearchSession(object):
    """Class encapsulating a typeahead search session."""

    def __init__(self):
        self.trie = TypeaheadSearchTrie()
        self.entries = {}

    def run_command(self, command):
        """Validate and execute a search command."""
        if command.startswith('ADD '):
            self.add(command[4:])
        elif command.startswith('DEL '):
            self.delete(command[4:])
        elif command.startswith('QUERY '):
            self.query(command[6:])
        elif command.startswith('WQUERY '):
            self.wquery(command[7:])
        else:
            raise ValueError(
                "Command \"{}\" is not of type ADD, DEL, QUERY,"
                " or WQUERY.".format(command)
            )

    def add(self, command):
        """Add a new item."""
        new_entry = Entry(*command.split(maxsplit=3))
        self.entries[new_entry.id] = new_entry

        for word in new_entry.data.lower().split():
            # Attempt to strip punctuation off of each word before storing
            # it as a search token.
            self.trie.add(word.strip(string.punctuation), new_entry)

    def delete(self, command):
        """Delete an item."""
        del self.entries[command]

    def query(self, command):
        """Perform a search."""

    def wquery(self, command):
        """Perform a weighted search."""


def main(session=None):
    """Main search loop."""
    if not session:
        session = TypeAheadSearchSession()

    # Get the number of expected commands.
    num_commands = sys.stdin.readline().strip()

    # Fetch each command from the input.
    for i in range(num_commands):
        command = sys.stdin.readline().strip()
        session.run_command(command)


if __name__ == '__main__':
    main()

@@ -1582,24 +1582,845 @@
 a search.%22%22%22
+%0A        num_results, search_words = command.split(maxsplit=1)%0A        num_results = int(num_results)%0A        search_words = search_words.split()%0A%0A        # Get the results set for one of the search words.%0A        results = self.trie.search(%0A            search_words.pop().strip(string.punctuation).lower()%0A        )%0A%0A        # Intersect the results sets for the remaining search words into%0A        # the first results set.%0A        for word in search_words.strip(string.punctuation).lower():%0A            results &= self.trie.search(%0A                word.strip(string.punctuation).lower()%0A            )%0A%0A        output = %5B%5D%0A        for i, entry in enumerate(sorted(results, key=lambda e: e.score)):%0A            if i %3E= num_results:%0A                break%0A%0A            output.append(entry.id)%0A%0A        return ''.join(output)
 %0A%0A    def wq

Implement search.