#!/usr/bin/env python
"""Message registry for apitools."""

import collections
import contextlib
import json

from protorpc import descriptor
from protorpc import messages
import six

from apitools.gen import extended_descriptor
from apitools.gen import util

TypeInfo = collections.namedtuple('TypeInfo', ('type_name', 'variant'))


class MessageRegistry(object):

    """Registry for message types.

    This closely mirrors a messages.FileDescriptor, but adds additional
    attributes (such as message and field descriptions) and some extra
    code for validation and cycle detection.
    """

    # Type information from these two maps comes from here:
    #  https://developers.google.com/discovery/v1/type-format
    PRIMITIVE_TYPE_INFO_MAP = {
        'string': TypeInfo(type_name='string',
                           variant=messages.StringField.DEFAULT_VARIANT),
        'integer': TypeInfo(type_name='integer',
                            variant=messages.IntegerField.DEFAULT_VARIANT),
        'boolean': TypeInfo(type_name='boolean',
                            variant=messages.BooleanField.DEFAULT_VARIANT),
        'number': TypeInfo(type_name='number',
                           variant=messages.FloatField.DEFAULT_VARIANT),
        'any': TypeInfo(type_name='extra_types.JsonValue',
                        variant=messages.Variant.MESSAGE),
    }

    PRIMITIVE_FORMAT_MAP = {
        'int32': TypeInfo(type_name='integer',
                          variant=messages.Variant.INT32),
        'uint32': TypeInfo(type_name='integer',
                           variant=messages.Variant.UINT32),
        'int64': TypeInfo(type_name='string',
                          variant=messages.Variant.INT64),
        'uint64': TypeInfo(type_name='string',
                           variant=messages.Variant.UINT64),
        'double': TypeInfo(type_name='number',
                           variant=messages.Variant.DOUBLE),
        'float': TypeInfo(type_name='number',
                          variant=messages.Variant.FLOAT),
        'byte': TypeInfo(type_name='byte',
                         variant=messages.BytesField.DEFAULT_VARIANT),
        'date': TypeInfo(type_name='extra_types.DateField',
                         variant=messages.Variant.STRING),
        'date-time': TypeInfo(
            type_name='protorpc.message_types.DateTimeMessage',
            variant=messages.Variant.MESSAGE),
    }

    def __init__(self, client_info, names, description,
                 root_package_dir, base_files_package):
        self.__names = names
        self.__client_info = client_info
        self.__package = client_info.package
        self.__description = util.CleanDescription(description)
        self.__root_package_dir = root_package_dir
        self.__base_files_package = base_files_package
        self.__file_descriptor = extended_descriptor.ExtendedFileDescriptor(
            package=self.__package, description=self.__description)
        # Add required imports
        self.__file_descriptor.additional_imports = [
            'from protorpc import messages as _messages',
        ]
        # Map from scoped names (i.e. Foo.Bar) to MessageDescriptors.
        self.__message_registry = collections.OrderedDict()
        # A set of types that we're currently adding (for cycle detection).
        self.__nascent_types = set()
        # A set of types for which we've seen a reference but no
        # definition; if this set is nonempty, validation fails.
        self.__unknown_types = set()
        # Used for tracking paths during message creation
        self.__current_path = []
        # Where to register created messages
        self.__current_env = self.__file_descriptor
        # TODO(craigcitro): Add a `Finalize` method.

    @property
    def file_descriptor(self):
        self.Validate()
        return self.__file_descriptor

    def WriteProtoFile(self, printer):
        """Write the messages file to out as proto."""
        self.Validate()
        extended_descriptor.WriteMessagesFile(
            self.__file_descriptor, self.__package, self.__client_info.version,
            printer)

    def WriteFile(self, printer):
        """Write the messages file to out."""
        self.Validate()
        extended_descriptor.WritePythonFile(
            self.__file_descriptor, self.__package, self.__client_info.version,
            printer)

    def Validate(self):
        mysteries = self.__nascent_types or self.__unknown_types
        if mysteries:
            raise ValueError('Malformed MessageRegistry: %s' % mysteries)

    def __ComputeFullName(self, name):
        return '.'.join(map(six.text_type, self.__current_path[:] + [name]))

    def __AddImport(self, new_import):
        if new_import not in self.__file_descriptor.additional_imports:
            self.__file_descriptor.additional_imports.append(new_import)

    def __DeclareDescriptor(self, name):
        self.__nascent_types.add(self.__ComputeFullName(name))

    def __RegisterDescriptor(self, new_descriptor):
        """Register the given descriptor in this registry."""
        if not isinstance(new_descriptor, (
                extended_descriptor.ExtendedMessageDescriptor,
                extended_descriptor.ExtendedEnumDescriptor)):
            raise ValueError('Cannot add descriptor of type %s' % (
                type(new_descriptor),))
        full_name = self.__ComputeFullName(new_descriptor.name)
        if full_name in self.__message_registry:
            raise ValueError(
                'Attempt to re-register descriptor %s' % full_name)
        if full_name not in self.__nascent_types:
            raise ValueError('Directly adding types is not supported')
        new_descriptor.full_name = full_name
        self.__message_registry[full_name] = new_descriptor
        if isinstance(new_descriptor,
                      extended_descriptor.ExtendedMessageDescriptor):
            self.__current_env.message_types.append(new_descriptor)
        elif isinstance(new_descriptor,
                        extended_descriptor.ExtendedEnumDescriptor):
            self.__current_env.enum_types.append(new_descriptor)
        self.__unknown_types.discard(full_name)
        self.__nascent_types.remove(full_name)

    def LookupDescriptor(self, name):
        return self.__GetDescriptorByName(name)

    def LookupDescriptorOrDie(self, name):
        message_descriptor = self.LookupDescriptor(name)
        if message_descriptor is None:
            raise ValueError('No message descriptor named "%s"', name)
        return message_descriptor

    def __GetDescriptor(self, name):
        return self.__GetDescriptorByName(self.__ComputeFullName(name))

    def __GetDescriptorByName(self, name):
        if name in self.__message_registry:
            return self.__message_registry[name]
        if name in self.__nascent_types:
            raise ValueError(
                'Cannot retrieve type currently being created: %s' % name)
        return None

    @contextlib.contextmanager
    def __DescriptorEnv(self, message_descriptor):
        # TODO(craigcitro): Typecheck?
        previous_env = self.__current_env
        self.__current_path.append(message_descriptor.name)
        self.__current_env = message_descriptor
        yield
        self.__current_path.pop()
        self.__current_env = previous_env

    def AddEnumDescriptor(self, name, description,
                          enum_values, enum_descriptions):
        """Add a new EnumDescriptor named name with the given enum values."""
        message = extended_descriptor.ExtendedEnumDescriptor()
        message.name = self.__names.ClassName(name)
        message.description = util.CleanDescription(description)
        self.__DeclareDescriptor(message.name)
        for index, (enum_name, enum_description) in enumerate(
                zip(enum_values, enum_descriptions)):
            enum_value = extended_descriptor.ExtendedEnumValueDescriptor()
            enum_value.name = self.__names.NormalizeEnumName(enum_name)
            if enum_value.name != enum_name:
                message.enum_mappings.append(
                    extended_descriptor.ExtendedEnumDescriptor.JsonEnumMapping(
                        python_name=enum_value.name, json_name=enum_name))
                self.__AddImport('from %s import encoding' %
                                 self.__base_files_package)
            enum_value.number = index
            enum_value.description = util.CleanDescription(
                enum_description or '<no description>')
            message.values.append(enum_value)
        self.__RegisterDescriptor(message)

    def __DeclareMessageAlias(self, schema, alias_for):
        """Declare schema as an alias for alias_for."""
        # TODO(craigcitro): This is a hack. Remove it.
        message = extended_descriptor.ExtendedMessageDescriptor()
        message.name = self.__names.ClassName(schema['id'])
        message.alias_for = alias_for
        self.__DeclareDescriptor(message.name)
        self.__AddImport('from %s import extra_types' %
                         self.__base_files_package)
        self.__RegisterDescriptor(message)

    def __AddAdditionalProperties(self, message, schema, properties):
        """Add an additionalProperties field to message."""
        additional_properties_info = schema['additionalProperties']
        entries_type_name = self.__AddAdditionalPropertyType(
            message.name, additional_properties_info)
        description = util.CleanDescription(
            additional_properties_info.get('description'))
        if description is None:
            description = 'Additional properties of type %s' % message.name
        attrs = {
            'items': {
                '$ref': entries_type_name,
            },
            'description': description,
            'type': 'array',
        }
        field_name = 'additionalProperties'
        message.fields.append(self.__FieldDescriptorFromProperties(
            field_name, len(properties) + 1, attrs))
        self.__AddImport('from %s import encoding' % self.__base_files_package)
        message.decorators.append(
            'encoding.MapUnrecognizedFields(%r)' % field_name)

    def AddDescriptorFromSchema(self, schema_name, schema):
        """Add a new MessageDescriptor named schema_name based on schema."""
        # TODO(craigcitro): Is schema_name redundant?
        if self.__GetDescriptor(schema_name):
            return
        if schema.get('enum'):
            self.__DeclareEnum(schema_name, schema)
            return
        if schema.get('type') == 'any':
            self.__DeclareMessageAlias(schema, 'extra_types.JsonValue')
            return
        if schema.get('type') != 'object':
            raise ValueError('Cannot create message descriptors for type %s',
                             schema.get('type'))
        message = extended_descriptor.ExtendedMessageDescriptor()
        message.name = self.__names.ClassName(schema['id'])
        message.description = util.CleanDescription(schema.get(
            'description', 'A %s object.' % message.name))
        self.__DeclareDescriptor(message.name)
        with self.__DescriptorEnv(message):
            properties = schema.get('properties', {})
            for index, (name, attrs) in enumerate(sorted(properties.items())):
                field = self.__FieldDescriptorFromProperties(
                    name, index + 1, attrs)
                message.fields.append(field)
                if field.name != name:
                    message.field_mappings.append(
                        type(message).JsonFieldMapping(
                            python_name=field.name, json_name=name))
                    self.__AddImport(
                        'from %s import encoding' % self.__base_files_package)
            if 'additionalProperties' in schema:
                self.__AddAdditionalProperties(message, schema, properties)
        self.__RegisterDescriptor(message)

    def __AddAdditionalPropertyType(self, name, property_schema):
        """Add a new nested AdditionalProperty message."""
        new_type_name = 'AdditionalProperty'
        property_schema = dict(property_schema)
        # We drop the description here on purpose, so the resulting
        # messages are less repetitive.
        property_schema.pop('description', None)
        description = 'An additional property for a %s object.' % name
        schema = {
            'id': new_type_name,
            'type': 'object',
            'description': description,
            'properties': {
                'key': {
                    'type': 'string',
                    'description': 'Name of the additional property.',
                },
                'value': property_schema,
            },
        }
        self.AddDescriptorFromSchema(new_type_name, schema)
        return new_type_name

    def __AddEntryType(self, entry_type_name, entry_schema, parent_name):
        """Add a type for a list entry."""
        entry_schema.pop('description', None)
        description = 'Single entry in a %s.' % parent_name
        schema = {
            'id': entry_type_name,
            'type': 'object',
            'description': description,
            'properties': {
                'entry': {
                    'type': 'array',
                    'items': entry_schema,
                },
            },
        }
        self.AddDescriptorFromSchema(entry_type_name, schema)
        return entry_type_name

    def __FieldDescriptorFromProperties(self, name, index, attrs):
        """Create a field descriptor for these attrs."""
        field = descriptor.FieldDescriptor()
        field.name = self.__names.CleanName(name)
        field.number = index
        field.label = self.__ComputeLabel(attrs)
        new_type_name_hint = self.__names.ClassName(
            '%sValue' % self.__names.ClassName(name))
        type_info = self.__GetTypeInfo(attrs, new_type_name_hint)
        field.type_name = type_info.type_name
        field.variant = type_info.variant
        if 'default' in attrs:
            # TODO(craigcitro): Correctly handle non-primitive default values.
            default = attrs['default']
            if not (field.type_name == 'string' or
                    field.variant == messages.Variant.ENUM):
                default = str(json.loads(default))
            if field.variant == messages.Variant.ENUM:
                default = self.__names.NormalizeEnumName(default)
            field.default_value = default
        extended_field = extended_descriptor.ExtendedFieldDescriptor()
        extended_field.name = field.name
        extended_field.description = util.CleanDescription(
            attrs.get('description', 'A %s attribute.' % field.type_name))
        extended_field.field_descriptor = field
        return extended_field

    @staticmethod
    def __ComputeLabel(attrs):
        if attrs.get('required', False):
            return descriptor.FieldDescriptor.Label.REQUIRED
        elif attrs.get('type') == 'array':
            return descriptor.FieldDescriptor.Label.REPEATED
        elif attrs.get('repeated'):
            return descriptor.FieldDescriptor.Label.REPEATED
        return descriptor.FieldDescriptor.Label.OPTIONAL

    def __DeclareEnum(self, enum_name, attrs):
        description = util.CleanDescription(attrs.get('description', ''))
        enum_values = attrs['enum']
        enum_descriptions = attrs.get(
            'enumDescriptions', [''] * len(enum_values))
        self.AddEnumDescriptor(enum_name, description,
                               enum_values, enum_descriptions)
        self.__AddIfUnknown(enum_name)
        return TypeInfo(type_name=enum_name, variant=messages.Variant.ENUM)

    def __AddIfUnknown(self, type_name):
        type_name = self.__names.ClassName(type_name)
        full_type_name = self.__ComputeFullName(type_name)
        if (full_type_name not in self.__message_registry.keys() and
                type_name not in self.__message_registry.keys()):
            self.__unknown_types.add(type_name)

    def __GetTypeInfo(self, attrs, name_hint):
        """Return a TypeInfo object for attrs, creating one if needed."""

        type_ref = self.__names.ClassName(attrs.get('$ref'))
        type_name = attrs.get('type')
        if not (type_ref or type_name):
            raise ValueError('No type found for %s' % attrs)

        if type_ref:
            self.__AddIfUnknown(type_ref)
            # We don't actually know this is a message -- it might be an
            # enum. However, we can't check that until we've created all the
            # types, so we come back and fix this up later.
            return TypeInfo(
                type_name=type_ref, variant=messages.Variant.MESSAGE)

        if 'enum' in attrs:
            enum_name = '%sValuesEnum' % name_hint
            return self.__DeclareEnum(enum_name, attrs)

        if 'format' in attrs:
            type_info = self.PRIMITIVE_FORMAT_MAP.get(attrs['format'])
            if type_info is None:
                # If we don't recognize the format, the spec says we fall back
                # to just using the type name.
                if type_name in self.PRIMITIVE_TYPE_INFO_MAP:
                    return self.PRIMITIVE_TYPE_INFO_MAP[type_name]
                raise ValueError('Unknown type/format "%s"/"%s"' % (
                    attrs['format'], type_name))
            if (type_info.type_name.startswith('protorpc.message_types.') or
                    type_info.type_name.startswith('message_types.')):
                self.__AddImport(
                    'from protorpc import message_types as _message_types')
            if type_info.type_name.startswith('extra_types.'):
                self.__AddImport(
                    'from %s import extra_types' % self.__base_files_package)
            return type_info

        if type_name in self.PRIMITIVE_TYPE_INFO_MAP:
            type_info = self.PRIMITIVE_TYPE_INFO_MAP[type_name]
            return type_info

        if type_name == 'array':
            items = attrs.get('items')
            if not items:
                raise ValueError('Array type with no item type: %s' % attrs)
            entry_name_hint = self.__names.ClassName(
                items.get('title') or '%sListEntry' % name_hint)
            entry_label = self.__ComputeLabel(items)
            if entry_label == descriptor.FieldDescriptor.Label.REPEATED:
                parent_name = self.__names.ClassName(
                    items.get('title') or name_hint)
                entry_type_name = self.__AddEntryType(
                    entry_name_hint, items.get('items'), parent_name)
                return TypeInfo(type_name=entry_type_name,
                                variant=messages.Variant.MESSAGE)
            else:
                return self.__GetTypeInfo(items, entry_name_hint)
        elif type_name == 'any':
            self.__AddImport('from %s import extra_types' %
                             self.__base_files_package)
            return self.PRIMITIVE_TYPE_INFO_MAP['any']
        elif type_name == 'object':
            # TODO(craigcitro): Think of a better way to come up with names.
            if not name_hint:
                raise ValueError(
                    'Cannot create subtype without some name hint')
            schema = dict(attrs)
            schema['id'] = name_hint
            self.AddDescriptorFromSchema(name_hint, schema)
            self.__AddIfUnknown(name_hint)
            return TypeInfo(
                type_name=name_hint, variant=messages.Variant.MESSAGE)

        raise ValueError('Unknown type: %s' % type_name)

    def FixupMessageFields(self):
        for message_type in self.file_descriptor.message_types:
            self._FixupMessage(message_type)

    def _FixupMessage(self, message_type):
        with self.__DescriptorEnv(message_type):
            for field in message_type.fields:
                if field.field_descriptor.variant == messages.Variant.MESSAGE:
                    field_type_name = field.field_descriptor.type_name
                    field_type = self.LookupDescriptor(field_type_name)
                    if isinstance(field_type,
                                  extended_descriptor.ExtendedEnumDescriptor):
                        field.field_descriptor.variant = messages.Variant.ENUM
            for submessage_type in message_type.message_types:
                self._FixupMessage(submessage_type)