text:
"""Student progress trackers."""

__author__ = 'Sean Lip (sll@google.com)'

import datetime
import logging
import os

import courses
import transforms

from models import QuestionDAO
from models import QuestionGroupDAO
from models import StudentPropertyEntity
from tools import verify

# Names of component tags that are tracked for progress calculations.
TRACKABLE_COMPONENTS = [
    'question',
    'question-group',
]

class UnitLessonCompletionTracker(object):
    """Tracks student completion for a unit/lesson-based linear course."""

    PROPERTY_KEY = 'linear-course-completion'

    # Here are representative examples of the keys for the various entities
    # used in this class:
    #   Unit 1: u.1
    #   Unit 1, Lesson 1: u.1.l.1
    #   Unit 1, Lesson 1, Activity 0: u.1.l.1.a.0
    #   Unit 1, Lesson 1, Activity 0, Block 4: u.1.l.1.a.0.b.4
    #   Assessment 'Pre': s.Pre
    # At the moment, we do not divide assessments into blocks.
    #
    # The following keys were added in v1.5:
    #   Unit 1, Lesson 1, HTML: u.1.l.1.h.0
    #   Unit 1, Lesson 1, HTML, Component with instanceid id: u.1.l.1.h.0.c.id
    #
    # The number after the 'h' and 'a' codes is always zero, since a lesson may
    # have at most one HTML body and one activity.
    #
    # IMPORTANT NOTE: The values of the keys mean different things depending on
    # whether the entity is a composite entity or not.
    # If it is a composite entity (unit, lesson, activity), then the value is
    #   - 0 if none of its sub-entities has been completed
    #   - 1 if some, but not all, of its sub-entities have been completed
    #   - 2 if all its sub-entities have been completed.
    # If it is not a composite entity (i.e. block, assessment, component), then
    # the value is just the number of times the event has been triggered.

    # Constants for recording the state of composite entities.
    # TODO(sll): Change these to enums.
    NOT_STARTED_STATE = 0
    IN_PROGRESS_STATE = 1
    COMPLETED_STATE = 2

    MULTIPLE_CHOICE = 'multiple choice'
    MULTIPLE_CHOICE_GROUP = 'multiple choice group'
    QUESTION_GROUP = 'question-group'
    QUESTION = 'question'

    EVENT_CODE_MAPPING = {
        'course': 'r',
        'course_forced': 'r',
        'unit': 'u',
        'unit_forced': 'u',
        'lesson': 'l',
        'activity': 'a',
        'html': 'h',
        'block': 'b',
        'assessment': 's',
        'component': 'c',
    }
    COMPOSITE_ENTITIES = [
        EVENT_CODE_MAPPING['course'],
        EVENT_CODE_MAPPING['unit'],
        EVENT_CODE_MAPPING['lesson'],
        EVENT_CODE_MAPPING['activity'],
        EVENT_CODE_MAPPING['html']
    ]

    def __init__(self, course):
        self._course = course

    def _get_course(self):
        return self._course

    def get_activity_as_python(self, unit_id, lesson_id):
        """Gets the corresponding activity as a Python object."""
        root_name = 'activity'
        course = self._get_course()
        activity_text = course.app_context.fs.get(
            os.path.join(course.app_context.get_home(),
                         course.get_activity_filename(unit_id, lesson_id)))
        content, noverify_text = verify.convert_javascript_to_python(
            activity_text, root_name)
        activity = verify.evaluate_python_expression_from_text(
            content, root_name, verify.Activity().scope, noverify_text)
        return activity

    def _get_course_key(self):
        return '%s.0' % (
            self.EVENT_CODE_MAPPING['course'],
        )

    def _get_unit_key(self, unit_id):
        return '%s.%s' % (self.EVENT_CODE_MAPPING['unit'], unit_id)

    def _get_lesson_key(self, unit_id, lesson_id):
        return '%s.%s.%s.%s' % (
            self.EVENT_CODE_MAPPING['unit'], unit_id,
            self.EVENT_CODE_MAPPING['lesson'], lesson_id
        )

    def _get_activity_key(self, unit_id, lesson_id):
        return '%s.%s.%s.%s.%s.%s' % (
            self.EVENT_CODE_MAPPING['unit'], unit_id,
            self.EVENT_CODE_MAPPING['lesson'], lesson_id,
            self.EVENT_CODE_MAPPING['activity'], 0
        )

    def _get_html_key(self, unit_id, lesson_id):
        return '%s.%s.%s.%s.%s.%s' % (
            self.EVENT_CODE_MAPPING['unit'], unit_id,
            self.EVENT_CODE_MAPPING['lesson'], lesson_id,
            self.EVENT_CODE_MAPPING['html'], 0
        )

    def _get_component_key(self, unit_id, lesson_id, component_id):
        return '%s.%s.%s.%s.%s.%s.%s.%s' % (
            self.EVENT_CODE_MAPPING['unit'], unit_id,
            self.EVENT_CODE_MAPPING['lesson'], lesson_id,
            self.EVENT_CODE_MAPPING['html'], 0,
            self.EVENT_CODE_MAPPING['component'], component_id
        )

    def _get_block_key(self, unit_id, lesson_id, block_id):
        return '%s.%s.%s.%s.%s.%s.%s.%s' % (
            self.EVENT_CODE_MAPPING['unit'], unit_id,
            self.EVENT_CODE_MAPPING['lesson'], lesson_id,
            self.EVENT_CODE_MAPPING['activity'], 0,
            self.EVENT_CODE_MAPPING['block'], block_id
        )

    def _get_assessment_key(self, assessment_id):
        assessment_key = '%s.%s' % (
            self.EVENT_CODE_MAPPING['assessment'], assessment_id)

        # If this assessment is used as a "lesson" within a unit, prepend
        # the unit identifier.
        parent_unit = self._get_course().get_parent_unit(assessment_id)
        if parent_unit:
            assessment_key = '.'.join([self._get_unit_key(parent_unit.unit_id),
                                       assessment_key])
        return assessment_key

    def get_entity_type_from_key(self, progress_entity_key):
        return progress_entity_key.split('.')[-2]

    def determine_if_composite_entity(self, progress_entity_key):
        return self.get_entity_type_from_key(
            progress_entity_key) in self.COMPOSITE_ENTITIES

    def get_valid_component_ids(self, unit_id, lesson_id):
        """Returns a list of cpt ids representing trackable components."""
        components = []
        for cpt_name in TRACKABLE_COMPONENTS:
            all_cpts = self._get_course().get_components_with_name(
                unit_id, lesson_id, cpt_name)
            components += [
                cpt['instanceid'] for cpt in all_cpts if cpt['instanceid']]
        return components

    def get_valid_block_ids(self, unit_id, lesson_id):
        """Returns a list of block ids representing interactive activities."""
        valid_blocks_data = self._get_valid_blocks_data(unit_id, lesson_id)
        return [block[0] for block in valid_blocks_data]

    def get_valid_blocks(self, unit_id, lesson_id):
        """Returns a list of blocks representing interactive activities."""
        valid_blocks_data = self._get_valid_blocks_data(unit_id, lesson_id)
        return [block[1] for block in valid_blocks_data]

    def _get_valid_blocks_data(self, unit_id, lesson_id):
        """Returns a list of (b_id, block) representing trackable activities."""
        valid_blocks = []

        # Check if activity exists before calling get_activity_as_python.
        unit = self._get_course().find_unit_by_id(unit_id)
        lesson = self._get_course().find_lesson_by_id(unit, lesson_id)
        if unit and lesson and lesson.activity:
            # Get the activity corresponding to this unit/lesson combination.
            activity = self.get_activity_as_python(unit_id, lesson_id)
            for block_id in range(len(activity['activity'])):
                block = activity['activity'][block_id]
                if isinstance(block, dict):
                    valid_blocks.append((block_id, block))
        return valid_blocks

    def get_id_to_questions_dict(self):
        """Returns a dict that maps each question to a list of its answers.

        Returns:
            A dict that represents the questions in lessons. The keys of this
            dict are question ids, and the corresponding values are dicts, each
            containing the following five key-value pairs:
            - answers: a list of 0's with length corresponding to number of
                choices a question has.
            - location: str. href value of the location of the question in the
                course.
            - num_attempts: int. Number of attempts for this question. This is
                used as the denominator when calculating the average score for a
                question. This value may differ from the sum of the elements in
                'answers' because of event entities that record an answer but
                not a score.
            - score: int. Aggregated value of the scores.
            - label: str. Human readable identifier for this question.
        """
        id_to_questions = {}
        for unit in self._get_course().get_units_of_type(verify.UNIT_TYPE_UNIT):
            unit_id = unit.unit_id
            for lesson in self._get_course().get_lessons(unit_id):
                lesson_id = lesson.lesson_id
                # Add mapping dicts for questions in old-style activities.
                if lesson.activity:
                    blocks = self._get_valid_blocks_data(unit_id, lesson_id)
                    for block_index, (block_id, block) in enumerate(blocks):
                        if block['questionType'] == self.MULTIPLE_CHOICE:
                            # Old style question.
                            id_to_questions.update(
                                self._create_old_style_question_dict(
                                    block, block_id, block_index, unit, lesson))

                        elif (block['questionType'] ==
                              self.MULTIPLE_CHOICE_GROUP):
                            # Old style multiple choice group.
                            for ind, q in enumerate(block['questionsList']):
                                id_to_questions.update(
                                    self._create_old_style_question_dict(
                                        q, block_id, block_index, unit,
                                        lesson, index=ind))

                # Add mapping dicts for CBv1.5 style questions.
                if lesson.objectives:
                    for cpt in self._get_course().get_question_components(
                            unit_id, lesson_id):
                        # CB v1.5 style questions.
                        id_to_questions.update(
                            self._create_v15_lesson_question_dict(
                                cpt, unit, lesson))

                    for cpt in self._get_course().get_question_group_components(
                            unit_id, lesson_id):
                        # CB v1.5 style question groups.
                        id_to_questions.update(
                            self._create_v15_lesson_question_group_dict(
                                cpt, unit, lesson))

        return id_to_questions

    def get_id_to_assessments_dict(self):
        """Returns a dict that maps each question to a list of its answers.

        Returns:
            A dict that represents the questions in assessments. The keys of
            this dict are question ids, and the corresponding values are dicts,
            each containing the following five key-value pairs:
            - answers: a list of 0's with length corresponding to number of
                choices a question has.
            - location: str. href value of the location of the question in the
                course.
            - num_attempts: int. Number of attempts for this question. This is
                used as the denominator when calculating the average score for a
                question. This value may differ from the sum of the elements in
                'answers' because of event entities that record an answer but
                not a score.
            - score: int. Aggregated value of the scores.
            - label: str. Human readable identifier for this question.
        """
        id_to_assessments = {}
        for assessment in self._get_course().get_assessment_list():
            if not self._get_course().needs_human_grader(assessment):
                assessment_components = self._get_course(
                    ).get_assessment_components(assessment.unit_id)
                # CB v1.5 style assessments.
                for cpt in assessment_components:
                    if cpt['cpt_name'] == self.QUESTION_GROUP:
                        id_to_assessments.update(
                            self._create_v15_assessment_question_group_dict(
                                cpt, assessment))
                    elif cpt['cpt_name'] == self.QUESTION:
                        id_to_assessments.update(
                            self._create_v15_assessment_question_dict(
                                cpt, assessment))

                # Old style javascript assessments.
                try:
                    content = self._get_course().get_assessment_content(
                        assessment)
                    id_to_assessments.update(
                        self._create_old_style_assessment_dict(
                            content['assessment'], assessment))
                except AttributeError:
                    # Assessment file does not exist.
                    continue

        return id_to_assessments

    def _get_link_for_assessment(self, assessment_id):
        return 'assessment?name=%s' % (assessment_id)

    def _get_link_for_activity(self, unit_id, lesson_id):
        return 'activity?unit=%s&lesson=%s' % (unit_id, lesson_id)

    def _get_link_for_lesson(self, unit_id, lesson_id):
        return 'unit?unit=%s&lesson=%s' % (unit_id, lesson_id)

    def _create_v15_question_dict(self, q_id, label, link, num_choices):
        """Returns a dict that represents CB v1.5 style question."""
        return {
            q_id: {
                'answer_counts': [0] * num_choices,
                'label': label,
                'location': link,
                'score': 0,
                'num_attempts': 0
            }
        }

    def _create_v15_lesson_question_dict(self, cpt, unit, lesson):
        try:
            question = QuestionDAO.load(cpt['quid'])
            if question.type == question.MULTIPLE_CHOICE:
                q_id = 'u.%s.l.%s.c.%s' % (
                    unit.unit_id, lesson.lesson_id, cpt['instanceid'])
                label = 'Unit %s Lesson %s, Question %s' % (
                    unit.index, lesson.index, question.description)
                link = self._get_link_for_lesson(unit.unit_id, lesson.lesson_id)
                num_choices = len(question.dict['choices'])
                return self._create_v15_question_dict(
                    q_id, label, link, num_choices)
            else:
                return {}
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, cpt)
            return {}

    def _create_v15_lesson_question_group_dict(self, cpt, unit, lesson):
        try:
            question_group = QuestionGroupDAO.load(cpt['qgid'])
            questions = {}
            for ind, quid in enumerate(question_group.question_ids):
                question = QuestionDAO.load(quid)
                if question.type == question.MULTIPLE_CHOICE:
                    q_id = 'u.%s.l.%s.c.%s.i.%s' % (
                        unit.unit_id, lesson.lesson_id, cpt['instanceid'], ind)
                    label = ('Unit %s Lesson %s, Question Group %s Question %s'
                             % (unit.index, lesson.index,
                                question_group.description,
                                question.description))
                    link = self._get_link_for_lesson(
                        unit.unit_id, lesson.lesson_id)
                    num_choices = len(question.dict['choices'])
                    questions.update(self._create_v15_question_dict(
                        q_id, label, link, num_choices))
            return questions
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, cpt)
            return {}

    def _create_v15_assessment_question_group_dict(self, cpt, assessment):
        try:
            question_group = QuestionGroupDAO.load(cpt['qgid'])
            questions = {}
            for ind, quid in enumerate(question_group.question_ids):
                question = QuestionDAO.load(quid)
                if question.type == question.MULTIPLE_CHOICE:
                    q_id = 's.%s.c.%s.i.%s' % (
                        assessment.unit_id, cpt['instanceid'], ind)
                    label = '%s, Question Group %s Question %s' % (
                        assessment.title, question_group.description,
                        question.description)
                    link = self._get_link_for_assessment(assessment.unit_id)
                    num_choices = len(question.dict['choices'])
                    questions.update(
                        self._create_v15_question_dict(
                            q_id, label, link, num_choices))
            return questions
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, cpt)
            return {}

    def _create_v15_assessment_question_dict(self, cpt, assessment):
        try:
            question = QuestionDAO.load(cpt['quid'])
            if question.type == question.MULTIPLE_CHOICE:
                q_id = 's.%s.c.%s' % (assessment.unit_id, cpt['instanceid'])
                label = '%s, Question %s' % (
                    assessment.title, question.description)
                link = self._get_link_for_assessment(assessment.unit_id)
                num_choices = len(question.dict['choices'])
                return self._create_v15_question_dict(
                    q_id, label, link, num_choices)
            else:
                return {}
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, cpt)
            return {}

    def _create_old_style_question_dict(self, block, block_id, block_index,
                                        unit, lesson, index=None):
        try:
            if index is not None:
                # Question is in a multiple choice group.
                b_id = 'u.%s.l.%s.b.%s.i.%s' % (
                    unit.unit_id, lesson.lesson_id, block_id, index)
                label = 'Unit %s Lesson %s Activity, Item %s Part %s' % (
                    unit.index, lesson.index, block_index + 1, index + 1)
            else:
                b_id = 'u.%s.l.%s.b.%s' % (
                    unit.unit_id, lesson.lesson_id, block_id)
                label = 'Unit %s Lesson %s Activity, Item %s' % (
                    unit.index, lesson.index, block_index + 1)
            return {
                b_id: {
                    'answer_counts': [0] * len(block['choices']),
                    'label': label,
                    'location': self._get_link_for_activity(
                        unit.unit_id, lesson.lesson_id),
                    'score': 0,
                    'num_attempts': 0
                }
            }
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, block)
            return {}

    def _create_old_style_assessment_dict(self, content, assessment):
        try:
            questions = {}
            for ind, question in enumerate(content['questionsList']):
                if 'choices' in question:
                    questions.update(
                        {
                            's.%s.i.%s' % (assessment.unit_id, ind): {
                                'answer_counts': [0] * len(question['choices']),
                                'label': '%s, Question %s' % (
                                    assessment.title, ind + 1),
                                'location': self._get_link_for_assessment(
                                    assessment.unit_id),
                                'score': 0,
                                'num_attempts': 0
                            }
                        }
                    )
            return questions
        except Exception as e:  # pylint: disable-msg=broad-except
            logging.error(
                'Failed to process the question data. '
                'Error: %s, data: %s', e, content)
            return {}

    def _update_course(self, progress, student):
        event_key = self._get_course_key()
        if self._get_entity_value(progress, event_key) == self.COMPLETED_STATE:
            return

        self._set_entity_value(progress, event_key, self.IN_PROGRESS_STATE)
        course = self._get_course()
        for unit in course.get_track_matching_student(student):
            if course.get_parent_unit(unit.unit_id):
                # Completion of an assessment-as-lesson rolls up to its
                # containing unit; it is not considered for overall course
                # completion (except insofar as assessment completion
                # contributes to the completion of its owning unit)
                pass
            else:
                if unit.type == verify.UNIT_TYPE_ASSESSMENT:
                    if not self.is_assessment_completed(progress, unit.unit_id):
                        return
                elif unit.type == verify.UNIT_TYPE_UNIT:
                    unit_state = self.get_unit_status(progress, unit.unit_id)
                    if unit_state != self.COMPLETED_STATE:
                        return
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_course_forced(self, progress):
        """Force state of course to completed."""

        event_key = self._get_course_key()
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_unit(self, progress, event_key):
        """Updates a unit's progress if all its lessons have been completed."""
        split_event_key = event_key.split('.')
        assert len(split_event_key) == 2
        unit_id = split_event_key[1]

        if self._get_entity_value(progress, event_key) == self.COMPLETED_STATE:
            return

        # Record that at least one lesson in this unit has been completed.
        self._set_entity_value(progress, event_key, self.IN_PROGRESS_STATE)

        # Check if all lessons in this unit have been completed.
        lessons = self._get_course().get_lessons(unit_id)
        for lesson in lessons:
            if (self.get_lesson_status(
                    progress,
                    unit_id, lesson.lesson_id) != self.COMPLETED_STATE):
                return

        # Check whether pre/post assessments in this unit have been completed.
        unit = self._get_course().find_unit_by_id(unit_id)
        pre_assessment_id = unit.pre_assessment
        if (pre_assessment_id and
            not self.get_assessment_status(progress, pre_assessment_id)):
            return
        post_assessment_id = unit.post_assessment
        if (post_assessment_id and
            not self.get_assessment_status(progress, post_assessment_id)):
            return

        # Record that all lessons in this unit have been completed.
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_unit_forced(self, progress, event_key):
        """Force-mark a unit as completed, ignoring normal criteria."""

        # Record that all lessons in this unit have been completed.
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_lesson(self, progress, event_key):
        """Updates a lesson's progress based on the progress of its children."""
        split_event_key = event_key.split('.')
        assert len(split_event_key) == 4
        unit_id = split_event_key[1]
        lesson_id = split_event_key[3]

        if self._get_entity_value(progress, event_key) == self.COMPLETED_STATE:
            return

        # Record that at least one part of this lesson has been completed.
        self._set_entity_value(progress, event_key, self.IN_PROGRESS_STATE)

        lessons = self._get_course().get_lessons(unit_id)
        for lesson in lessons:
            if str(lesson.lesson_id) == lesson_id and lesson:
                # Is the activity completed?
                if (lesson.activity and self.get_activity_status(
                        progress, unit_id, lesson_id) != self.COMPLETED_STATE):
                    return

                # Are all components of the lesson completed?
                if (self.get_html_status(
                        progress, unit_id, lesson_id) != self.COMPLETED_STATE):
                    return

        # Record that all activities in this lesson have been completed.
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_activity(self, progress, event_key):
        """Updates activity's progress when all interactive blocks are done."""
        split_event_key = event_key.split('.')
        assert len(split_event_key) == 6
        unit_id = split_event_key[1]
        lesson_id = split_event_key[3]

        if self._get_entity_value(progress, event_key) == self.COMPLETED_STATE:
            return

        # Record that at least one block in this activity has been completed.
        self._set_entity_value(progress, event_key, self.IN_PROGRESS_STATE)

        valid_block_ids = self.get_valid_block_ids(unit_id, lesson_id)
        for block_id in valid_block_ids:
            if not self.is_block_completed(
                    progress, unit_id, lesson_id, block_id):
                return

        # Record that all blocks in this activity have been completed.
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    def _update_html(self, progress, event_key):
        """Updates html's progress when all interactive blocks are done."""
        split_event_key = event_key.split('.')
        assert len(split_event_key) == 6
        unit_id = split_event_key[1]
        lesson_id = split_event_key[3]

        if self._get_entity_value(progress, event_key) == self.COMPLETED_STATE:
            return

        # Record that at least one block in this activity has been completed.
        self._set_entity_value(progress, event_key, self.IN_PROGRESS_STATE)

        cpt_ids = self.get_valid_component_ids(unit_id, lesson_id)
        for cpt_id in cpt_ids:
            if not self.is_component_completed(
                    progress, unit_id, lesson_id, cpt_id):
                return

        # Record that all blocks in this activity have been completed.
        self._set_entity_value(progress, event_key, self.COMPLETED_STATE)

    UPDATER_MAPPING = {
        'activity': _update_activity,
        'course': _update_course,
        'course_forced': _update_course_forced,
        'html': _update_html,
        'lesson': _update_lesson,
        'unit': _update_unit,
        'unit_forced': _update_unit_forced,
    }

    # Dependencies for recording derived events. The key is the current
    # event, and the value is a tuple, each element of which contains:
    # - the dependent entity to be updated
    # - the transformation to apply to the id of the current event to get the
    #       id for the derived parent event
    DERIVED_EVENTS = {
        'block': (
            {
                'entity': 'activity',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
        'activity': (
            {
                'entity': 'lesson',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
        'lesson': (
            {
                'entity': 'unit',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
        'component': (
            {
                'entity': 'html',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
        'html': (
            {
                'entity': 'lesson',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
        'assessment': (
            {
                'entity': 'unit',
                'generate_parent_id': (lambda s: '.'.join(s.split('.')[:-2]))
            },
        ),
    }

    def force_course_completed(self, student):
        self._put_event(
            student, 'course_forced', self._get_course_key())

    def force_unit_completed(self, student, unit_id):
        """Records that the given student has completed a unit.

        NOTE: This should not generally be used directly.  The definition
        of completing a unit is generally taken to be completion of all
        parts of all components of a unit (assessments, lessons,
        activities in lessons, etc.  Directly marking a unit as complete
        is provided only for manual marking where the student feels "done",
        but has not taken a fully completionist approach to the material.

        Args:
          student: A logged-in, registered student object.
          unit_id: The ID of the unit to be marked as complete.
        """
        self._put_event(
            student, 'unit_forced', self._get_unit_key(unit_id))

    def put_activity_completed(self, student, unit_id, lesson_id):
        """Records that the given student has completed an activity."""
        if not self._get_course().is_valid_unit_lesson_id(unit_id, lesson_id):
            return
        self._put_event(
            student, 'activity', self._get_activity_key(unit_id, lesson_id))

    def put_html_completed(self, student, unit_id, lesson_id):
        """Records that the given student has completed a lesson page."""
        if not self._get_course().is_valid_unit_lesson_id(unit_id, lesson_id):
            return
        self._put_event(
            student, 'html', self._get_html_key(unit_id, lesson_id))

    def put_block_completed(self, student, unit_id, lesson_id, block_id):
        """Records that the given student has completed an activity block."""
        if not self._get_course().is_valid_unit_lesson_id(unit_id, lesson_id):
            return
        if block_id not in self.get_valid_block_ids(unit_id, lesson_id):
            return
        self._put_event(
            student,
            'block',
            self._get_block_key(unit_id, lesson_id, block_id)
        )

    def put_component_completed(self, student, unit_id, lesson_id, cpt_id):
        """Records completion of a component in a lesson body."""
        if not self._get_course().is_valid_unit_lesson_id(unit_id, lesson_id):
            return
        if cpt_id not in self.get_valid_component_ids(unit_id, lesson_id):
            return
        self._put_event(
            student,
            'component',
            self._get_component_key(unit_id, lesson_id, cpt_id)
        )

    def put_assessment_completed(self, student, assessment_id):
        """Records that the given student has completed the given assessment."""
        if not self._get_course().is_valid_assessment_id(assessment_id):
            return
        self._put_event(
            student, 'assessment', self._get_assessment_key(assessment_id))

    def put_activity_accessed(self, student, unit_id, lesson_id):
        """Records that the given student has accessed this activity."""
        # This method currently exists because we need to mark activities
        # without interactive blocks as 'completed' when they are accessed.
        if not self.get_valid_block_ids(unit_id, lesson_id):
            self.put_activity_completed(student, unit_id, lesson_id)

    def put_html_accessed(self, student, unit_id, lesson_id):
        """Records that the given student has accessed this lesson page."""
        # This method currently exists because we need to mark lesson bodies
        # without interactive blocks as 'completed' when they are accessed.
        if not self.get_valid_component_ids(unit_id, lesson_id):
            self.put_html_completed(student, unit_id, lesson_id)

    def _put_event(self, student, event_entity, event_key):
        """Starts a cascade of updates in response to an event taking place."""
        if student.is_transient or event_entity not in self.EVENT_CODE_MAPPING:
            return

        progress = self.get_or_create_progress(student)

        self._update_event(
            student, progress, event_entity, event_key, direct_update=True)

        progress.updated_on = datetime.datetime.now()
        progress.put()

    def _update_event(self, student, progress, event_entity, event_key,
                      direct_update=False):
        """Updates statistics for the given event, and for derived events.

        Args:
          student: the student
          progress: the StudentProgressEntity for the student
          event_entity: the name of the affected entity (unit, lesson, etc.)
          event_key: the key for the recorded event
          direct_update: True if this event is being updated explicitly; False
              if it is being auto-updated.
        """
        if direct_update or event_entity not in self.UPDATER_MAPPING:
            if event_entity in self.UPDATER_MAPPING:
                # This is a derived event, so directly mark it as completed.
                self._set_entity_value(
                    progress, event_key, self.COMPLETED_STATE)
            else:
                # This is not a derived event, so increment its counter by one.
                self._inc(progress, event_key)
        else:
            self.UPDATER_MAPPING[event_entity](self, progress, event_key)

        if event_entity in self.DERIVED_EVENTS:
            for derived_event in self.DERIVED_EVENTS[event_entity]:
                parent_event_key = derived_event['generate_parent_id'](
                    event_key)
                if parent_event_key:
                    # Event entities may contribute upwards to more than one
                    # kind of container.  Only pass the notification up to the
                    # handler that our event_key indicates we actually have.
                    leaf_type = self.get_entity_type_from_key(parent_event_key)
                    event_entity = derived_event['entity']
                    if leaf_type == self.EVENT_CODE_MAPPING[event_entity]:
                        self._update_event(
                            student=student,
                            progress=progress,
                            event_entity=event_entity,
                            event_key=parent_event_key)
                else:
                    # Only update course status when we are at the top of
                    # a containment list
                    self._update_course(progress, student)
        else:
            # Or only update course status when we are doing something not
            # in derived events (Unit, typically).
            self._update_course(progress, student)

    def get_course_status(self, progress):
        return self._get_entity_value(progress, self._get_course_key())

    def get_unit_status(self, progress, unit_id):
        return self._get_entity_value(progress, self._get_unit_key(unit_id))

    def get_lesson_status(self, progress, unit_id, lesson_id):
        return self._get_entity_value(
            progress, self._get_lesson_key(unit_id, lesson_id))

    def get_activity_status(self, progress, unit_id, lesson_id):
        return self._get_entity_value(
            progress, self._get_activity_key(unit_id, lesson_id))

    def get_html_status(self, progress, unit_id, lesson_id):
        return self._get_entity_value(
            progress, self._get_html_key(unit_id, lesson_id))

    def get_block_status(self, progress, unit_id, lesson_id, block_id):
        return self._get_entity_value(
            progress, self._get_block_key(unit_id, lesson_id, block_id))

    def get_assessment_status(self, progress, assessment_id):
        return self._get_entity_value(
            progress, self._get_assessment_key(assessment_id))

    def is_block_completed(self, progress, unit_id, lesson_id, block_id):
        value = self._get_entity_value(
            progress, self._get_block_key(unit_id, lesson_id, block_id))
        return value is not None and value > 0

    def is_component_completed(self, progress, unit_id, lesson_id, cpt_id):
        value = self._get_entity_value(
            progress, self._get_component_key(unit_id, lesson_id, cpt_id))
        return value is not None and value > 0

    def is_assessment_completed(self, progress, assessment_id):
        value = self._get_entity_value(
            progress, self._get_assessment_key(assessment_id))
        return value is not None and value > 0

    @classmethod
    def get_or_create_progress(cls, student):
        progress = StudentPropertyEntity.get(student, cls.PROPERTY_KEY)
        if not progress:
            progress = StudentPropertyEntity.create(
                student=student, property_name=cls.PROPERTY_KEY)
            progress.put()
        return progress

    def get_course_progress(self, student):
        """Return [NOT_STARTED|IN_PROGRESS|COMPLETED]_STATE for course."""
        progress = self.get_or_create_progress(student)
        return self.get_course_status(progress) or self.NOT_STARTED_STATE

    def get_unit_progress(self, student, progress=None):
        """Returns a dict with the states of each unit."""
        if student.is_transient:
            return {}

        units = self._get_course().get_units()
        if progress is None:
            progress = self.get_or_create_progress(student)

        result = {}
        for unit in units:
            if unit.type == verify.UNIT_TYPE_ASSESSMENT:
                result[unit.unit_id] = self.is_assessment_completed(
                    progress, unit.unit_id)
            elif unit.type == verify.UNIT_TYPE_UNIT:
                value = self.get_unit_status(progress, unit.unit_id)
                result[unit.unit_id] = value or 0

        return result

    def get_unit_percent_complete(self, student):
        """Returns a dict with each unit's completion in [0.0, 1.0]."""
        if student.is_transient:
            return {}

        course = self._get_course()
        units = course.get_units()
        assessment_scores = {int(s['id']): s['score'] / 100.0
                             for s in course.get_all_scores(student)}
        result = {}
        progress = self.get_or_create_progress(student)
        for unit in units:
            # Assessments are scored as themselves.
            if unit.type == verify.UNIT_TYPE_ASSESSMENT:
                result[unit.unit_id] = assessment_scores[unit.unit_id]
            elif unit.type == verify.UNIT_TYPE_UNIT:
                if (unit.pre_assessment and
                    assessment_scores[unit.pre_assessment] >= 1.0):
                    # Use pre-assessment iff it exists and student scored 100%
                    result[unit.unit_id] = 1.0
                else:
                    # Otherwise, count % completion on lessons within unit.
                    num_completed = 0
                    lesson_progress = self.get_lesson_progress(
                        student, unit.unit_id, progress=progress)
                    if not lesson_progress:
                        result[unit.unit_id] = 0.0
                    else:
                        for lesson in lesson_progress.values():
                            if lesson['has_activity']:
                                # Lessons that have activities must be
                                # activity-complete as well as HTML complete.
                                if (lesson['html'] == self.COMPLETED_STATE and
                                    lesson['activity'] == self.COMPLETED_STATE):
                                    num_completed += 1
                            else:
                                # Lessons without activities just need HTML
                                if lesson['html'] == self.COMPLETED_STATE:
                                    num_completed += 1
                        result[unit.unit_id] = round(
                            num_completed / float(len(lesson_progress)), 3)
        return result

    def get_lesson_progress(self, student, unit_id, progress=None):
        """Returns a dict saying which lessons in this unit are completed."""
        if student.is_transient:
            return {}

        lessons = self._get_course().get_lessons(unit_id)
        if progress is None:
            progress = self.get_or_create_progress(student)

        result = {}
        for lesson in lessons:
            result[lesson.lesson_id] = {
                'html': self.get_html_status(
                    progress, unit_id, lesson.lesson_id) or 0,
                'activity': self.get_activity_status(
                    progress, unit_id, lesson.lesson_id) or 0,
                'has_activity': lesson.has_activity,
            }
        return result

    def get_component_progress(self, student, unit_id, lesson_id, cpt_id):
        """Returns the progress status of the given component."""
        if student.is_transient:
            return 0

        progress = self.get_or_create_progress(student)
        return self.is_component_completed(
            progress, unit_id, lesson_id, cpt_id) or 0

    def _get_entity_value(self, progress, event_key):
        if not progress.value:
            return None
        return transforms.loads(progress.value).get(event_key)

    def _set_entity_value(self, student_property, key, value):
        """Sets the integer value of a student property.

        Note: this method does not commit the change. The calling method should
        call put() on the StudentPropertyEntity.

        Args:
          student_property: the StudentPropertyEntity
          key: the student property whose value should be incremented
          value: the value to increment this property by
        """
        try:
            progress_dict = transforms.loads(student_property.value)
        except (AttributeError, TypeError):
            progress_dict = {}

        progress_dict[key] = value
        student_property.value = transforms.dumps(progress_dict)

    def _inc(self, student_property, key, value=1):
        """Increments the integer value of a student property.

        Note: this method does not commit the change. The calling method should
        call put() on the StudentPropertyEntity.

        Args:
          student_property: the StudentPropertyEntity
          key: the student property whose value should be incremented
          value: the value to increment this property by
        """
        try:
            progress_dict = transforms.loads(student_property.value)
        except (AttributeError, TypeError):
            progress_dict = {}

        if key not in progress_dict:
            progress_dict[key] = 0

        progress_dict[key] += value
        student_property.value = transforms.dumps(progress_dict)

class ProgressStats(object):
    """Defines the course structure definition for course progress tracking."""

    def __init__(self, course):
        self._course = course
        self._tracker = UnitLessonCompletionTracker(course)

    def compute_entity_dict(self, entity, parent_ids):
        """Computes the course structure dictionary.

        Args:
            entity: str. Represents for which level of entity the dict is being
                computed. Valid entity levels are defined as keys to the dict
                defined below, COURSE_STRUCTURE_DICT.
            parent_ids: list of ids necessary to get children of the current
                entity.
        Returns:
            A nested dictionary representing the structure of the course.
            Every other level of the dictionary consists of a key, the label of
            the entity level defined by EVENT_CODE_MAPPING in
            UnitLessonCompletionTracker, whose value is a dictionary
            INSTANCES_DICT. The keys of INSTANCES_DICT are instance_ids of the
            corresponding entities, and the values are the entity_dicts of the
            instance's children, in addition to a field called 'label'. Label
            represents the user-facing name of the entity rather than
            its intrinsic id. If one of these values is empty, this means
            that the corresponding entity has no children.

            Ex:
            A Course with the following outlined structure:
                Pre Assessment
                Unit 1
                    Lesson 1
                Unit 2

            will have the following dictionary representation:
                {
                    's': {
                        1: {
                            'label': 'Pre Assessment'
                        }
                    },
                    'u': {
                        2: {
                            'l': {
                                3: {
                                    'label': 1
                                }
                            },
                            'label': 1
                        },
                        4: {
                            'label': 2
                        }
                    }
                    'label': 'UNTITLED COURSE'
                }
        """
        entity_dict = {'label': self._get_label(entity, parent_ids)}
        for child_entity, get_children_ids in self.COURSE_STRUCTURE_DICT[
                entity]['children']:
            child_entity_dict = {}
            for child_id in get_children_ids(self, *parent_ids):
                new_parent_ids = parent_ids + [child_id]
                child_entity_dict[child_id] = self.compute_entity_dict(
                    child_entity, new_parent_ids)
            entity_dict[UnitLessonCompletionTracker.EVENT_CODE_MAPPING[
                child_entity]] = child_entity_dict
        return entity_dict

    def _get_course(self):
        return self._course

    def _get_unit_ids_of_type_unit(self):
        units = self._get_course().get_units_of_type(verify.UNIT_TYPE_UNIT)
        return [unit.unit_id for unit in units]

    def _get_assessment_ids(self):
        assessments = self._get_course().get_assessment_list()
        return [a.unit_id for a in assessments]

    def _get_lesson_ids(self, unit_id):
        lessons = self._get_course().get_lessons(unit_id)
        return [lesson.lesson_id for lesson in lessons]

    def _get_activity_ids(self, unit_id, lesson_id):
        unit = self._get_course().find_unit_by_id(unit_id)
        if self._get_course().find_lesson_by_id(unit, lesson_id).activity:
            return [0]
        return []

    def _get_html_ids(self, unused_unit_id, unused_lesson_id):
        return [0]

    def _get_block_ids(self, unit_id, lesson_id, unused_activity_id):
        return self._tracker.get_valid_block_ids(unit_id, lesson_id)

    def _get_component_ids(self, unit_id, lesson_id, unused_html_id):
        return self._tracker.get_valid_component_ids(unit_id, lesson_id)

    def _get_label(self, entity, parent_ids):
        return self.ENTITY_TO_HUMAN_READABLE_NAME_DICT[entity](
            self, *parent_ids)

    def _get_course_label(self):
        # pylint: disable-msg=protected-access
        return courses.Course.get_environ(self._get_course().app_context)[
            'course']['title']

    def _get_unit_label(self, unit_id):
        unit = self._get_course().find_unit_by_id(unit_id)
        return 'Unit %s' % unit.index

    def _get_assessment_label(self, unit_id):
        assessment = self._get_course().find_unit_by_id(unit_id)
        return assessment.title

    def _get_lesson_label(self, unit_id, lesson_id):
        unit = self._get_course().find_unit_by_id(unit_id)
        lesson = self._get_course().find_lesson_by_id(unit, lesson_id)
        return lesson.index

    def _get_activity_label(self, unit_id, lesson_id, unused_activity_id):
        return str('L%s.%s' % (
            self._get_course().find_unit_by_id(unit_id).index,
            self._get_lesson_label(unit_id, lesson_id)))

    def _get_html_label(self, unit_id, lesson_id, unused_html_id):
        return self._get_activity_label(unit_id, lesson_id, unused_html_id)

    def _get_block_label(self, unit_id, lesson_id, unused_activity_id,
                         block_id):
        return str('L%s.%s.%s' % (
            self._get_course().find_unit_by_id(unit_id).index,
            self._get_lesson_label(unit_id, lesson_id),
            block_id))

    def _get_component_label(self, unit_id, lesson_id, unused_html_id,
                             component_id):
        return self._get_block_label(
            unit_id, lesson_id, unused_html_id, component_id)

    # Outlines the structure of the course. The key is the entity level, and
    # its value is a dictionary with following keys and its values:
    #   'children': list of tuples. Each tuple consists of string representation
    #               of the child entity(ex: 'lesson') and a function to get the
    #               children elements. If the entity does not have children, the
    #               value will be an empty list.
    #   'id': instance_id of the entity. If the entity is represented by a class
    #         with an id attribute(ex: units), string representation of the
    #         attribute is stored here. If the entity is defined by a dictionary
    #         (ex: components), then the value is the string 'None'.
    #
    COURSE_STRUCTURE_DICT = {
        'course': {
            'children': [('unit', _get_unit_ids_of_type_unit),
                         ('assessment', _get_assessment_ids)],
        },
        'unit': {
            'children': [('lesson', _get_lesson_ids)],
        },
        'assessment': {
            'children': [],
        },
        'lesson': {
            'children': [('activity', _get_activity_ids),
                         ('html', _get_html_ids)],
        },
        'activity': {
            'children': [('block', _get_block_ids)],
        },
        'html': {
            'children': [('component', _get_component_ids)],
        },
        'block': {
            'children': [],
        },
        'component': {
            'children': [],
        }
    }

    ENTITY_TO_HUMAN_READABLE_NAME_DICT = {
        'course': _get_course_label,
        'unit': _get_unit_label,
        'assessment': _get_assessment_label,
        'lesson': _get_lesson_label,
        'activity': _get_activity_label,
        'html': _get_html_label,
        'block': _get_block_label,
        'component': _get_component_label
    }


text:
require 'metaforce'
require 'metaforce/reporters/deploy_reporter'
require 'fileutils'

class SfOpticon::Environment < ActiveRecord::Base
  attr_reader :log, :config
  validates_uniqueness_of :name,
    :message => "This organization is already configured."
  attr_accessible :name,
    :host,
    :username,
    :password,
    :securitytoken,
    :production,
    :locked

  has_many :sf_objects
  has_one  :branch

  # Setup variables
  after_initialize do |env|
    @log = SfOpticon::Logger
    @config = SfOpticon::Settings.salesforce
  end

  ##
  # This method is called when an environment is first created. This allows
  # us to reach out and create the remote repository if needed, or the branch.
  # This will also clone the branch.
  after_create do
    # Validate salesforce credentials before creating anything remote.
    unless sforce.credentials_are_valid?
      raise "Could not login to Salesforce! Please verify credentials."
    end

    if production
      # If we're a production environment then we need to create the remote
      # repository
      SfOpticon::Scm.adapter.create_remote_repository(name)
    end

    create_branch(name: production ? 'master' : name)
    snapshot

    # We ignore the actual metadata objects in salesforce unless this is
    # production. We do this because we want them to start at the same logical
    # place, as though the non-production environment was refreshed even if it
    # wasn't.
    if production
      sforce.retrieve :manifest => sforce.manifest(sf_objects),
        :extract_to => branch.local_path

      branch.ignore_package_xml
      branch.add_changes
      branch.commit("Initial push of production code")
      branch.push
    end
  end

  ##
  # Each call to sforce will reconfigure the metaforce client.
  def sforce
    SfOpticon::Salesforce.new(self)
  end

  ##
  # Locks the environment to prevent the changeset tool from making any
  # changes. This is necessary so that integration branches will have
  # a pristine head to work from
  def lock
    self.locked = true
    save!
  end

  ##
  # Unlocks the environment.
  def unlock
    self.locked = false
    save!
  end

  ##
  # Rebases our branch from production
  def integrate(src_env)
    branch.integrate(src_env)
  end

  ##
  # Deploys code changes to this environment.
  #
  # @param sf_objects [String] The list of sf_objects to deploy.
  # @param destructive [Boolean] True if this is a destructive changeset
  def deploy(sf_objects, destructive = false)
    staging_dir = if destructive
      stage_destructive(sf_objects)
    else
      stage_artifacts(sf_objects)
    end
    log.info { "Deploying artifacts staged in #{staging_dir} to #{name}"}

    sforce.client.deploy(File.join(staging_dir,'src'), { :run_all_tests => false })
    .on_complete { |job| log.info job.result }
    .on_error {|job| log.error { "Error: #{job.result.inspect}" } }
    .on_poll {|job| log.info "Polling..." }
    .perform
  end

  ##
  # Creates a destructive staging area.
  # The directory must follow this layout
  # src/
  #   package.xml <- Must be an empty manifest
  #   destructiveChanges.xml
  def stage_destructive(sf_objects)
    staging_dir = Dir.mktmpdir
    src_dir = File.join(staging_dir, 'src')
    Dir.mkdir src_dir
    log.info { "Creating destructive staging area at #{staging_dir}" }

    File.open(File.join(src_dir, 'package.xml'), 'w') do |f|
      f.write(sforce.manifest([]).to_xml)
    end

    File.open(File.join(src_dir, 'destructiveChanges.xml'), 'w') do |f|
      f.write(sforce.manifest(sf_objects).to_xml)
    end

    staging_dir
  end

  ##
  # Creates the staging area for a set of deployment artifacts.
  # The directory from which we deploy must follow this layout:
  # src/
  #   package.xml
  #   classes
  #   etc.
  def stage_artifacts(sf_objects)
    staging_dir = Dir.mktmpdir
    log.info { "Staging artifacts info #{staging_dir} for deployment." }

    # Create layout
    src_dir = File.join(staging_dir, 'src')
    sf_objects.collect {|o| o.dirname }.uniq.each do |dir|
      FileUtils.mkdir_p(File.join(src_dir, dir))
    end

    # Create the package manifest
    File.open(File.join(src_dir, 'package.xml'), 'w') do |f|
      f.puts(sforce.manifest(sf_objects).to_xml)
    end

    # Copy the files into staging
    sf_objects.each do |o|
      FileUtils.cp(o.fileset, File.join(src_dir, o.dirname))
    end

    staging_dir
  end

  ##
  # Removes all sf_objects (via delete_all to avoid instantiation cost), the
  # local repo directory, and itself. This does *not* remove any remote repos!
  def remove
    # We skip the instantiation and go straight to single
    # statement deletion
    sf_objects.delete_all

    # And we delete our remote branch
    unless production
      branch.delete_remote_branch
    end

    # Discard the org contents.
    begin
      FileUtils.remove_dir("#{SfOpticon::Settings.scm.local_path}/#{name}")
    rescue Errno::ENOENT
      # We pass if the directory is already gone
    end

    delete
  end

  # Create's a clean snapshot of all SF metadata related to the
  # configured types.
  def snapshot
    ## Env has to have it's current sf_objects wiped out
    log.info { "Deleting all sfobjects for #{name}" }
    sf_objects.delete_all

    SfOpticon::SfObject.transaction do
      sforce.gather_metadata.each do |o|
        sf_objects << SfOpticon::SfObject.create(o)
      end
      save!
    end
  end

  # Generates a changeset from the latest snapshot and the current
  # metadata information from Salesforce. This changeset is then
  # committed to both the database and the repository.
  #
  # Returns the changeset
  def changeset
    curr_snap = sforce.gather_metadata
    change_queue = SfOpticon::Changes::Diff.snap_diff(sf_objects, curr_snap)
    if change_queue.all_changes.size == 0
      log.info { "No changes detected in #{name}" }
      return
    end

    # We now have an array of objects that have been deleted, renamed, added,
    # or modified in the correct order. We will replay these changes into the
    # local repository and the database.

    # First we have to generate a manifest of the additions and modifications
    # to retrieve those new objects
    mods = [change_queue.additions, change_queue.modifications].flatten.map {|x| x.sf_object }

    # Retrieve the changes into a temporary directory
    dir = Dir.mktmpdir("changeset")
    sforce.retrieve(:manifest => sforce.manifest(mods), :extract_to => dir)

    # Now we replay the changes into the repo and the database
    change_queue.apply_change_queue(dir, branch.local_path) do |change|
      log.info { "DIFF: #{change.change_type} - #{change.sf_object[:full_name]}" }

      commit_message = "#{change.change_type} - #{change.sf_object[:full_name]}\n\n"
      if change.change_type == :deletion
        commit_message += "#{change.sf_object[:file_name]} deleted"
      else
        change.sf_object.keys.each do |key|
          commit_message += "#{key.to_s.camelize}: #{change.sf_object[key]}\n"
        end
      end

      change.apply(dir, branch.local_path)
      branch.add_changes
      branch.commit(commit_message, change.sf_object[:last_modified_by_name])

      case change.change_type
      when :deletion
        sf_objects.find_by_sfobject_id(change.sf_object[:sfobject_id]).delete
      when :modification
        sf_objects
        .find_by_sfobject_id(change.sf_object[:sfobject_id])
        .clobber(change.sf_object)
      when :addition
        sf_objects << sf_objects.new(change.sf_object)
      end

      save!
    end

    branch.push
    FileUtils.remove_entry_secure(dir)

    log.info { "Complete." }
    change_queue
  end
end

redpajama_set_name: RedPajamaGithub


text:
package com.google.javascript.jscomp.serialization;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.truth.Truth.assertThat;

import com.google.auto.value.AutoValue;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.SetMultimap;
import com.google.common.truth.extensions.proto.ProtoSubject;
import com.google.common.truth.extensions.proto.ProtoTruth;
import com.google.javascript.jscomp.colors.Color;
import com.google.javascript.jscomp.colors.ColorId;
import com.google.javascript.jscomp.colors.DebugInfo;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

@RunWith(JUnit4.class)
public class ColorSerializerTest {

  @Test
  public void generateEmptyTypePool() {
    // Expect an empty TypePool when no colors are added.
    new Tester()
        .init()
        // don't add anything
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(
            TypePool.newBuilder()
                // DebugInfo is present, but empty
                .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
                .build());

    new Tester()
        .setSerializationMode(SerializationOptions.SKIP_DEBUG_INFO)
        .init()
        // don't add anything
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(TypePool.getDefaultInstance());
  }

  @Test
  public void onlyAddAxiomaticColors() {
    // These are the colors that are never actually serialized, even when requested.
    // They are always just assumed to be there when deserializing and to have fixed indices as
    // indicated by their index in TypePointers.OFFSET_TO_AXIOMATIC_COLOR.
    final List<TestColor> axiomaticTestColors =
        TypePointers.OFFSET_TO_AXIOMATIC_COLOR.stream()
            .map(this::getAxiomaticTestColor)
            .collect(Collectors.toList());

    // Expect an empty TypePool when only axiomatic colors are added.
    new Tester()
        .init()
        // add all the axiomatic colors, but only those
        .addColors(axiomaticTestColors)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(
            TypePool.newBuilder()
                // No `type` values, because axiomatic colors are never serialized.
                // DebugInfo is empty, but present.
                .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
                .build());
  }

  @Test
  public void addASuperTypeRelationship() {
    final TestColor superTypeColor =
        new TestObjectColorBuilder().setColorId("sup_type").setTrimmedPoolOffset(0).build();
    final TestColor subTypeColor =
        new TestObjectColorBuilder().setColorId("sub_type").setTrimmedPoolOffset(1).build();
    final TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(superTypeColor.getExpectedTypeProto())
            .addType(subTypeColor.getExpectedTypeProto())
            .addDisambiguationEdges(
                SubtypingEdge.newBuilder()
                    .setSubtype(subTypeColor.getExpectedTypePointer())
                    .setSupertype(superTypeColor.getExpectedTypePointer())
                    .build())
            // empty DebugInfo
            .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
            .build();

    new Tester()
        .init()
        .linkSubColorToSuperColor(subTypeColor, superTypeColor)
        .addColor(superTypeColor)
        .addColor(subTypeColor)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void implicitAdditionOfSuperType() {
    final TestColor subTypeColor =
        new TestObjectColorBuilder()
            .setColorId("sub_type")
            // The subtype will be serialized first, because it is explicitly added.
            .setTrimmedPoolOffset(0)
            .build();
    final TestColor superTypeColor =
        new TestObjectColorBuilder()
            .setColorId("sup_type")
            // The supertype will be serialized second, because it gets added implicitly
            // by its relationship to the subtype.
            .setTrimmedPoolOffset(1)
            .build();
    final TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(subTypeColor.getExpectedTypeProto())
            .addType(superTypeColor.getExpectedTypeProto())
            .addDisambiguationEdges(
                SubtypingEdge.newBuilder()
                    .setSubtype(subTypeColor.getExpectedTypePointer())
                    .setSupertype(superTypeColor.getExpectedTypePointer())
                    .build())
            // DebugInfo is present but empty.
            .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
            .build();

    new Tester()
        .init()
        .linkSubColorToSuperColor(subTypeColor, superTypeColor)
        // Only explicitly add the subtype
        .addColor(subTypeColor)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void simulatedClasses() {
    // Simulate Colors for this class:
    //
    // class Base {
    //   static baseStaticField;
    //   baseField;
    // }
    //
    // We will explicitly add the color representing `typeof Base`.
    // The colors for `Base.prototype` and the instances of `Base` will be added implicitly
    // in that order.
    final PooledString baseFieldPooledString = PooledString.create("baseField", 1);
    final TestColor baseClassInstanceTestColor =
        new TestObjectColorBuilder()
            .setColorId("Base.ins")
            // 0 => explicitly added base class constructor
            // 1 => explicitly added child class constructor
            // Then this instance color gets added implicitly
            .setTrimmedPoolOffset(2)
            .setTypeName("Base")
            .addOwnProperty(baseFieldPooledString)
            .build();
    final TestColor baseClassPrototypeTestColor =
        new TestObjectColorBuilder()
            .setColorId("Base.pro")
            // This is implicitly added after the instance color above.
            .setTrimmedPoolOffset(3)
            .setTypeName("typeof Base.prototype")
            .build();
    final PooledString baseStaticFieldPooledString = PooledString.create("baseStaticField", 0);
    final TestColor baseClassConstructorTestColor =
        new TestObjectColorBuilder()
            .setColorId("Base.con")
            // This is the first explicitly-added type.
            .setTrimmedPoolOffset(0)
            .setTypeName("typeof Base")
            .setConstructor(true)
            .addPrototypeTestColor(baseClassPrototypeTestColor)
            .addInstanceTestColor(baseClassInstanceTestColor)
            .addOwnProperty(baseStaticFieldPooledString)
            .build();

    // class Child extends Base {
    //   static childStaticField;
    //   childField;
    // }
    // We will explicitly add the color representing `typeof Child`.
    // The colors for `Child.prototype` and the instances of `Child` will be added implicitly
    // in that order.
    final PooledString childFieldPooledString = PooledString.create("childField", 3);
    final TestColor childClassInstanceTestColor =
        new TestObjectColorBuilder()
            .setColorId("Child.in")
            // 0 => explicitly added base class constructor color
            // 1 => explicitly added child class constructor color
            // 2 => implicitly added base class instance color
            // 3 => implicitly added base class prototype color
            // Then this instance color gets added implicitly
            .setTrimmedPoolOffset(4)
            .setTypeName("Child")
            .addOwnProperty(childFieldPooledString)
            .build();
    final TestColor childClassPrototypeTestColor =
        new TestObjectColorBuilder()
            .setColorId("Child.pr")
            // This is implicitly added after the instance color above.
            .setTrimmedPoolOffset(5)
            .setTypeName("typeof Child.prototype")
            .build();
    final PooledString childStaticFieldPooledString = PooledString.create("childStaticField", 2);
    final TestColor childClassConstructorTestColor =
        new TestObjectColorBuilder()
            .setColorId("Child.co")
            // This is the second explicitly-added type
            .setTrimmedPoolOffset(1)
            .setTypeName("typeof Child")
            .setConstructor(true)
            .addPrototypeTestColor(childClassPrototypeTestColor)
            .addInstanceTestColor(childClassInstanceTestColor)
            .addOwnProperty(childStaticFieldPooledString)
            .build();

    final TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(baseClassConstructorTestColor.getExpectedTypeProto())
            .addType(childClassConstructorTestColor.getExpectedTypeProto())
            .addType(baseClassInstanceTestColor.getExpectedTypeProto())
            .addType(baseClassPrototypeTestColor.getExpectedTypeProto())
            .addType(childClassInstanceTestColor.getExpectedTypeProto())
            .addType(childClassPrototypeTestColor.getExpectedTypeProto())
            .addDisambiguationEdges(
                SubtypingEdge.newBuilder()
                    .setSubtype(childClassConstructorTestColor.getExpectedTypePointer())
                    .setSupertype(baseClassConstructorTestColor.getExpectedTypePointer())
                    .build())
            .addDisambiguationEdges(
                SubtypingEdge.newBuilder()
                    .setSubtype(childClassInstanceTestColor.getExpectedTypePointer())
                    .setSupertype(baseClassInstanceTestColor.getExpectedTypePointer())
                    .build())
            .addDisambiguationEdges(
                SubtypingEdge.newBuilder()
                    .setSubtype(childClassPrototypeTestColor.getExpectedTypePointer())
                    .setSupertype(baseClassPrototypeTestColor.getExpectedTypePointer())
                    .build())
            // DebugInfo is present but empty.
            .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
            .build();

    new Tester()
        .init()
        .addPooledString(baseStaticFieldPooledString)
        .addPooledString(baseFieldPooledString)
        .addPooledString(childStaticFieldPooledString)
        .addPooledString(childFieldPooledString)
        // Only explicitly add the constructor types
        .addColor(baseClassConstructorTestColor)
        .addColor(childClassConstructorTestColor)
        .linkSubColorToSuperColor(childClassConstructorTestColor, baseClassConstructorTestColor)
        .linkSubColorToSuperColor(childClassInstanceTestColor, baseClassInstanceTestColor)
        .linkSubColorToSuperColor(childClassPrototypeTestColor, baseClassPrototypeTestColor)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void includeMismatchesAndTypenames() {
    TestColor testColor1 =
        new TestObjectColorBuilder()
            .setColorId("color001")
            .setTrimmedPoolOffset(0)
            .setTypeName("typeName1")
            .build();
    TestColor testColor2 =
        new TestObjectColorBuilder()
            .setColorId("color002")
            .setTrimmedPoolOffset(1)
            .setTypeName("typeName2")
            .build();
    TestMismatch testMismatch1 = TestMismatch.create("location1", testColor1, testColor2);

    TestColor testColor3 =
        new TestObjectColorBuilder()
            .setColorId("color003")
            .setTrimmedPoolOffset(2)
            .setTypeName("typeName3")
            .build();
    TestColor testColor4 =
        new TestObjectColorBuilder()
            .setColorId("color004")
            .setTrimmedPoolOffset(3)
            .setTypeName("typeName4")
            .build();
    TestMismatch testMismatch2 = TestMismatch.create("location2", testColor3, testColor4);

    final TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(testColor1.getExpectedTypeProto())
            .addType(testColor2.getExpectedTypeProto())
            .addType(testColor3.getExpectedTypeProto())
            .addType(testColor4.getExpectedTypeProto())
            .setDebugInfo(
                TypePool.DebugInfo.newBuilder()
                    .addMismatch(testMismatch1.getExpectedMismatch())
                    .addMismatch(testMismatch2.getExpectedMismatch())
                    .build())
            .build();

    new Tester()
        .init()
        .addColors(testColor1, testColor2, testColor3, testColor4)
        .addMismatch(testMismatch1)
        .addMismatch(testMismatch2)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void skipMismatchesAndTypenames() {
    TestColor testColor1 =
        new TestObjectColorBuilder()
            .setColorId("color001")
            .setTrimmedPoolOffset(0)
            .setTypeName("typeName1")
            .setIncludeDebugInfoInProto(false)
            .build();
    TestColor testColor2 =
        new TestObjectColorBuilder()
            .setColorId("color002")
            .setTrimmedPoolOffset(1)
            .setTypeName("typeName2")
            .setIncludeDebugInfoInProto(false)
            .build();
    TestMismatch testMismatch1 = TestMismatch.create("location1", testColor1, testColor2);

    TestColor testColor3 =
        new TestObjectColorBuilder()
            .setColorId("color003")
            .setTrimmedPoolOffset(2)
            .setTypeName("typeName3")
            .setIncludeDebugInfoInProto(false)
            .build();
    TestColor testColor4 =
        new TestObjectColorBuilder()
            .setColorId("color004")
            .setTrimmedPoolOffset(3)
            .setTypeName("typeName4")
            .setIncludeDebugInfoInProto(false)
            .build();
    TestMismatch testMismatch2 = TestMismatch.create("location2", testColor3, testColor4);

    final TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(testColor1.getExpectedTypeProto())
            .addType(testColor2.getExpectedTypeProto())
            .addType(testColor3.getExpectedTypeProto())
            .addType(testColor4.getExpectedTypeProto())
            // No debug info
            .build();

    new Tester()
        .setSerializationMode(SerializationOptions.SKIP_DEBUG_INFO)
        .init()
        .addColors(testColor1, testColor2, testColor3, testColor4)
        .addMismatch(testMismatch1)
        .addMismatch(testMismatch2)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void filterPropertyNames() {
    // Tell ColorSerializer to only keep property names beginning with "kept"
    final Predicate<String> propertyNameFilter = propName -> propName.startsWith("kept");

    // Build a type that has a property name that should be kept and one that should be
    // dropped.
    final PooledString keptPropertyName = PooledString.create("keptProperty", 0);
    final PooledString droppedPropertyName = PooledString.create("droppedProperty", 1);
    TestColor testColor =
        new TestObjectColorBuilder()
            .setColorId("color001")
            .setTrimmedPoolOffset(0)
            .addOwnProperty(keptPropertyName)
            .addOwnProperty(droppedPropertyName)
            .build();

    final TypePool actualTypePool =
        new Tester()
            .setPropertyNameFilter(propertyNameFilter)
            .init()
            .addPooledString(keptPropertyName)
            .addPooledString(droppedPropertyName)
            .addColor(testColor)
            .generateTypePool()
            .getTypePool();

    final List<Integer> ownPropertyList =
        actualTypePool.getType(0).getObject().getOwnPropertyList();
    assertThat(ownPropertyList).containsExactly(keptPropertyName.getPoolOffset());
  }

  @Test
  public void addUnionImplicitlyAddsMembers() {
    final TestColor testColor1 =
        new TestObjectColorBuilder()
            .setColorId("color001")
            // We will explicitly add this color first
            .setTrimmedPoolOffset(1)
            .setTypeName("testColor1")
            .build();
    final TestColor testColor2 =
        new TestObjectColorBuilder()
            .setColorId("color002")
            .setTrimmedPoolOffset(2)
            .setTypeName("testColor2")
            .build();
    final TestColor unionTestColor =
        new TestUnionColorBuilder()
            .setTrimmedPoolOffset(0)
            .addTestColor(testColor1)
            .addTestColor(testColor2)
            .build();

    TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(unionTestColor.getExpectedTypeProto())
            .addType(testColor1.getExpectedTypeProto())
            .addType(testColor2.getExpectedTypeProto())
            // empty DebugInfo
            .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
            .build();

    new Tester()
        .init()
        .addColor(unionTestColor)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  @Test
  public void avoidAddingDuplicateUnions() {
    final TestColor testColor1 =
        new TestObjectColorBuilder()
            .setColorId("color001")
            // We will explicitly add this color first
            .setTrimmedPoolOffset(1)
            .setTypeName("testColor1")
            .build();
    final TestColor testColor2 =
        new TestObjectColorBuilder()
            .setColorId("color002")
            .setTrimmedPoolOffset(2)
            .setTypeName("testColor2")
            .build();
    final TestColor unionTestColor =
        new TestUnionColorBuilder()
            .setTrimmedPoolOffset(0)
            .addTestColor(testColor1)
            .addTestColor(testColor2)
            .build();
    // Creating a second union containing the same types will create a distinct
    // Java object, but it will have the same ColorId.
    final TestColor duplicateUnionTestColor =
        new TestUnionColorBuilder()
            .setTrimmedPoolOffset(0) // should end up getting the same offset
            .addTestColor(testColor1)
            .addTestColor(testColor2)
            .build();

    TypePool expectedTypePool =
        TypePool.newBuilder()
            .addType(unionTestColor.getExpectedTypeProto())
            // duplicate Union color should not be serialized
            .addType(testColor1.getExpectedTypeProto())
            .addType(testColor2.getExpectedTypeProto())
            // empty DebugInfo
            .setDebugInfo(TypePool.DebugInfo.getDefaultInstance())
            .build();

    new Tester()
        .init()
        .addColor(unionTestColor)
        // Try to add the duplicate union color. It should be ignored.
        .addColor(duplicateUnionTestColor)
        .generateTypePool()
        .assertThatTypePool()
        .isEqualTo(expectedTypePool);
  }

  /**
   * Represents a string that is stored in a StringPool, recording both its value and its offset.
   */
  @AutoValue
  abstract static class PooledString {
    public abstract String getValue();

    public abstract int getPoolOffset();

    static PooledString create(String value, int poolOffset) {
      return new AutoValue_ColorSerializerTest_PooledString(value, poolOffset);
    }
  }

  /**
   * Builds objects to represent both an object Color to be added and the TypePointer and TypeProto
   * we expect to be generated for it.
   */
  private static class TestObjectColorBuilder {
    private final ArrayList<TestColor> instanceTestColors = new ArrayList<>();
    private final ArrayList<TestColor> prototypeTestColors = new ArrayList<>();
    private final ArrayList<PooledString> ownProperties = new ArrayList<>();

    private ColorId colorId;
    private int trimmedPoolOffset = -1;
    // An empty string indicates that the type has no name.
    private String typeName = "";
    private boolean isConstructor = false;
    private boolean includeDebugInfoInProto = true;

    public TestObjectColorBuilder setColorId(String colorIdString) {
      checkArgument(colorIdString.length() <= 8, "color ID string too long: %s", colorIdString);
      return setColorId(ColorId.fromAscii(colorIdString));
    }

    public TestObjectColorBuilder setColorId(ColorId colorId) {
      this.colorId = colorId;
      return this;
    }

    public TestObjectColorBuilder setTrimmedPoolOffset(int trimmedPoolOffset) {
      checkState(trimmedPoolOffset >= 0, "invalid trimmedPoolOffset: %s", trimmedPoolOffset);
      this.trimmedPoolOffset = trimmedPoolOffset;
      return this;
    }

    public TestObjectColorBuilder setConstructor(boolean constructor) {
      isConstructor = constructor;
      return this;
    }

    public TestObjectColorBuilder setTypeName(String typeName) {
      this.typeName = typeName;
      return this;
    }

    public TestObjectColorBuilder addInstanceTestColor(TestColor instanceTestColor) {
      this.instanceTestColors.add(instanceTestColor);
      return this;
    }

    public TestObjectColorBuilder addPrototypeTestColor(TestColor prototypeTestColor) {
      this.prototypeTestColors.add(prototypeTestColor);
      return this;
    }

    public TestObjectColorBuilder addOwnProperty(PooledString ownProperty) {
      ownProperties.add(ownProperty);
      return this;
    }

    public TestObjectColorBuilder setIncludeDebugInfoInProto(boolean includeDebugInfoInProto) {
      this.includeDebugInfoInProto = includeDebugInfoInProto;
      return this;
    }

    TestColor build() {
      checkState(trimmedPoolOffset >= 0, "call setTrimmedPoolOffset() first");
      // There's no testing benefit to varying these values
      final boolean propertiesKeepOriginalName = false;
      final boolean isInvalidating = false;
      final boolean isClosureAssert = false;

      final Color.Builder colorBuilder =
          Color.singleBuilder()
              .setId(colorId)
              .setInvalidating(isInvalidating)
              .setPropertiesKeepOriginalName(propertiesKeepOriginalName)
              .setConstructor(isConstructor)
              .setClosureAssert(isClosureAssert)
              .setInstanceColors(
                  instanceTestColors.stream().map(TestColor::getColor).collect(toImmutableSet()))
              .setPrototypes(
                  prototypeTestColors.stream().map(TestColor::getColor).collect(toImmutableSet()))
              .setOwnProperties(
                  ownProperties.stream().map(PooledString::getValue).collect(toImmutableSet()));
      final TypePointer typePointer =
          TypePointer.newBuilder()
              .setPoolOffset(TypePointers.untrimOffset(trimmedPoolOffset))
              .build();
      final TypeProto.Builder typeProtoBuilder = TypeProto.newBuilder();
      final ObjectTypeProto.Builder objectTypeProtoBuilder = typeProtoBuilder.getObjectBuilder();
      objectTypeProtoBuilder
          .setUuid(colorId.asByteString())
          .setIsInvalidating(isInvalidating)
          .setPropertiesKeepOriginalName(propertiesKeepOriginalName)
          .setMarkedConstructor(isConstructor)
          .setClosureAssert(isClosureAssert)
          .addAllInstanceType(
              instanceTestColors.stream()
                  .map(TestColor::getExpectedTypePointer)
                  .collect(Collectors.toList()))
          .addAllPrototype(
              prototypeTestColors.stream()
                  .map(TestColor::getExpectedTypePointer)
                  .collect(Collectors.toList()))
          .addAllOwnProperty(
              ownProperties.stream().map(PooledString::getPoolOffset).collect(Collectors.toList()));
      if (!typeName.isEmpty()) {
        colorBuilder.setDebugInfo(DebugInfo.builder().setCompositeTypename(typeName).build());
        if (includeDebugInfoInProto) {
          objectTypeProtoBuilder.getDebugInfoBuilder().addTypename(typeName);
        }
      }
      return TestColor.create(colorBuilder.build(), typeProtoBuilder.build(), typePointer);
    }
  }

  static class TestUnionColorBuilder {
    private final ArrayList<TestColor> memberTestColors = new ArrayList<>();

    private int trimmedPoolOffset = -1;

    public TestUnionColorBuilder setTrimmedPoolOffset(int trimmedPoolOffset) {
      checkState(trimmedPoolOffset >= 0, "invalid trimmedPoolOffset: %s", trimmedPoolOffset);
      this.trimmedPoolOffset = trimmedPoolOffset;
      return this;
    }

    TestUnionColorBuilder addTestColor(TestColor testColor) {
      memberTestColors.add(testColor);
      return this;
    }

    TestColor build() {
      checkState(trimmedPoolOffset >= 0, "call setTrimmedPoolOffset() first");
      final ImmutableSet<Color> memberColors =
          memberTestColors.stream().map(TestColor::getColor).collect(toImmutableSet());
      Color color = Color.createUnion(memberColors);

      final TypePointer typePointer =
          TypePointer.newBuilder()
              .setPoolOffset(TypePointers.untrimOffset(trimmedPoolOffset))
              .build();

      final List<TypePointer> memberTypePoiners =
          memberTestColors.stream()
              .map(TestColor::getExpectedTypePointer)
              .collect(Collectors.toList());
      final TypeProto.Builder typeProtoBuilder = TypeProto.newBuilder();
      typeProtoBuilder.getUnionBuilder().addAllUnionMember(memberTypePoiners);
      final TypeProto typeProto = typeProtoBuilder.build();
      return TestColor.create(color, typeProto, typePointer);
    }
  }

  TestColor getAxiomaticTestColor(Color axiomaticColor) {
    final int poolOffset = TypePointers.OFFSET_TO_AXIOMATIC_COLOR.indexOf(axiomaticColor);
    checkArgument(poolOffset >= 0, "Not an axiomatic color: %s", axiomaticColor);
    return TestColor.create(
        axiomaticColor, null, TypePointer.newBuilder().setPoolOffset(poolOffset).build());
  }

  /** Represents a Color that has been or will be added to the ColorSerializer. */
  @AutoValue
  abstract static class TestColor {
    // The Color that will be added to the ColorSerializer.
    public abstract Color getColor();

    /**
     * The TypeProto we expect ColorSerializer to create for this Color.
     *
     * <p>For an axiomatic color this will return `null`, since those are never stored into a
     * `TypeProto`. Generally test code should call `getExpectedTypeProto()` instead in order to get
     * an exception if an attempt is made to serialize an axiomatic color.
     */
    @Nullable
    public abstract TypeProto getNullableExpectedTypeProto();

    // The TypePointer we expect ColorSerializer to create for this Color.
    public abstract TypePointer getExpectedTypePointer();

    public TypeProto getExpectedTypeProto() {
      return checkNotNull(getNullableExpectedTypeProto());
    }

    static TestColor create(
        Color color, TypeProto expectedTypeProto, TypePointer nullableExpectedTypePointer) {
      return new AutoValue_ColorSerializerTest_TestColor(
          color, expectedTypeProto, nullableExpectedTypePointer);
    }
  }

  /**
   * Represents a color/type mismatch both as it would appear in a ColorRegistry and in a TypePool.
   */
  @AutoValue
  abstract static class TestMismatch {
    public abstract String getLocationString();

    public abstract ImmutableList<TestColor> getTestColors();

    public List<Color> getColors() {
      return getTestColors().stream().map(TestColor::getColor).collect(Collectors.toList());
    }

    public TypePool.DebugInfo.Mismatch getExpectedMismatch() {
      final List<TypePointer> involvedColorTypePointers =
          getTestColors().stream()
              .map(TestColor::getExpectedTypePointer)
              .collect(Collectors.toList());
      return TypePool.DebugInfo.Mismatch.newBuilder()
          .setSourceRef(getLocationString())
          .addAllInvolvedColor(involvedColorTypePointers)
          .build();
    }

    /**
     * @param locationString Used as the `source_ref` string in the Mismatch proto
     * @param testColors Used as the `involved_color` repeated field in the Mismatch proto, in the
     *     order
     */
    static TestMismatch create(String locationString, TestColor... testColors) {
      return new AutoValue_ColorSerializerTest_TestMismatch(
          locationString, ImmutableList.copyOf(testColors));
    }
  }

  private static class Tester {
    // Simulates ColorRegistry#getDisambiguationSuperTypes()
    private final SetMultimap<Color, Color> subColorToSuperColorsMap = HashMultimap.create();
    // Simulates ColorRegistry#getMismatchLocationsForDebugging()
    private final SetMultimap<Color, String> colorToMismatchLocationStringsMap =
        HashMultimap.create();
    // Simulates StringPool.Builder#put()
    private final HashMap<String, Integer> stringToPoolOffsetMap = new HashMap<>();
    // Test object
    private ColorSerializer colorSerializer;
    // Default to the serialization mode that includes the most information.
    // If the test case includes information that doesn't make it into the expected result,
    // then I want the test case to clearly request that the information be omitted.
    private SerializationOptions serializationMode =
        SerializationOptions.INCLUDE_DEBUG_INFO_AND_EXPENSIVE_VALIDITY_CHECKS;
    // Tells ColorSerializer which property names to serialize
    private Predicate<String> propertyNameFilter = propName -> true;

    Tester() {}

    Tester setSerializationMode(SerializationOptions serializationMode) {
      checkState(colorSerializer == null, "call this method before init()");
      this.serializationMode = serializationMode;
      return this;
    }

    Tester setPropertyNameFilter(Predicate<String> propertyNameFilter) {
      checkState(colorSerializer == null, "call this method before init()");
      this.propertyNameFilter = propertyNameFilter;
      return this;
    }

    Tester init() {
      checkState(colorSerializer == null, "init() called more than once");
      colorSerializer =
          new ColorSerializer(
              serializationMode,
              str -> checkNotNull(stringToPoolOffsetMap.get(str), "unexpected string: %s", str),
              propertyNameFilter);
      return this;
    }

    /** Be prepared to respond to a request for this string. */
    Tester addPooledString(PooledString pooledString) {
      final String string = pooledString.getValue();
      checkState(!stringToPoolOffsetMap.containsKey(string), "duplicate string added: %s", string);
      stringToPoolOffsetMap.put(string, pooledString.getPoolOffset());
      return this;
    }

    Tester linkSubColorToSuperColor(Color subColor, Color superColor) {
      subColorToSuperColorsMap.put(subColor, superColor);
      return this;
    }

    Tester linkSubColorToSuperColor(TestColor subColor, TestColor superColor) {
      return linkSubColorToSuperColor(subColor.getColor(), superColor.getColor());
    }

    Tester addMismatch(TestMismatch testMismatch) {
      final String locationString = testMismatch.getLocationString();
      for (Color color : testMismatch.getColors()) {
        colorToMismatchLocationStringsMap.put(color, locationString);
      }
      return this;
    }

    Tester addColor(TestColor testColor) {
      final TypePointer typePointer = colorSerializer.addColor(testColor.getColor());
      ProtoTruth.assertThat(typePointer).isEqualTo(testColor.getExpectedTypePointer());
      return this;
    }

    Tester addColors(TestColor... testColors) {
      return addColors(Arrays.asList(testColors));
    }

    Tester addColors(List<TestColor> testColorList) {
      checkNotNull(colorSerializer, "call init() first");
      final List<Color> colors =
          testColorList.stream().map(TestColor::getColor).collect(Collectors.toList());
      final ImmutableList<TypePointer> typePointers = colorSerializer.addColors(colors);
      for (int i = 0; i < testColorList.size(); ++i) {
        ProtoTruth.assertThat(typePointers.get(i))
            .isEqualTo(testColorList.get(i).getExpectedTypePointer());
      }
      return this;
    }

    GenerateTypePoolTestResult generateTypePool() {
      checkNotNull(colorSerializer, "call init() first");
      final TypePool typePool =
          colorSerializer.generateTypePool(
              subtypeColor -> ImmutableSet.copyOf(subColorToSuperColorsMap.get(subtypeColor)),
              color -> ImmutableSet.copyOf(colorToMismatchLocationStringsMap.get(color)));
      return GenerateTypePoolTestResult.create(typePool);
    }
  }

  /** The result of Tester::generateTypePool() */
  @AutoValue
  abstract static class GenerateTypePoolTestResult {
    public abstract TypePool getTypePool();

    ProtoSubject assertThatTypePool() {
      return ProtoTruth.assertThat(getTypePool());
    }

    static GenerateTypePoolTestResult create(TypePool typePool) {
      return new AutoValue_ColorSerializerTest_GenerateTypePoolTestResult(typePool);
    }
  }
}