"""
Automatically generates Avro Keys for Kafka topics keys using the property 'partition-field'.
"""
from collections import OrderedDict
from typing import Dict, List, Any

from avro_preprocessor.avro_domain import Avro
from avro_preprocessor.avro_naming import AvroNaming
from avro_preprocessor.preprocessor_module import PreprocessorModule
from avro_preprocessor.schemas_container import SchemasContainer

__author__ = "Nicola Bova"
__copyright__ = "Copyright 2018, Jaumo GmbH"
__email__ = "nicola.bova@jaumo.com"


class KeysGenerator(PreprocessorModule):
    """
    Generates Avro Keys for all topics.
    """

    def __init__(self, schemas: SchemasContainer):
        super().__init__(schemas)
        self.schema_keys: Dict[str, Any] = {}

        self.admissible_types: List[str] = ['int', 'long', 'string', 'bytes']

    def process(self) -> None:
        """Process all schemas."""

        for name, schema in self.processed_schemas_iter():
            if self.schemas.paths.is_event_resource(name):
                key_name = AvroNaming.get_key_fully_qualified_name(name)

                if key_name not in self.schemas.autogenerated_keys:
                    # one key per topic

                    self.schema_keys[name] = []
                    self.traverse_schema(self.find_partition_keys_records, schema)
                    self.schema_keys[name] = sorted(self.schema_keys[name],
                                                    key=lambda k: k[Avro.Name])  # type: ignore

                    if self.schema_keys[name]:
                        avro_key = OrderedDict((
                            (Avro.Namespace, schema[Avro.Namespace]),
                            (Avro.Name, 'Key'),
                            (Avro.Doc, '{} {}'.format(
                                AvroNaming.autogenerated_key_doc,
                                AvroNaming.get_topic(self.schemas.paths.base_namespace, name))),
                            (Avro.Type, Avro.Record),
                            (Avro.Fields, self.schema_keys[name])
                        ))

                        self.schemas.autogenerated_keys[key_name] = avro_key

    def find_partition_keys_records(
            self, record: Avro.Node) -> None:
        """
        Finds property 'partition-field' inside schemas and saves the record they belong to.
        :param record: The schema
        """
        if isinstance(record, OrderedDict) and Avro.PartitionKey in record:
            if record[Avro.PartitionKey]:

                # type is one of admissible types or a logicalType of admissible types
                if record[Avro.Type] in self.admissible_types or (
                        # type is a logicalType
                        isinstance(record[Avro.Type], OrderedDict)
                        and Avro.Type in record[Avro.Type]
                        and Avro.LogicalType in record[Avro.Type]
                        and record[Avro.Type][Avro.Type] in self.admissible_types
                ):
                    if record not in self.schema_keys[self.current_schema_name]:
                        self.schema_keys[self.current_schema_name].append(record)
                else:
                    raise ValueError(
                        'Partition key in field {} of schema {} of type {}: '
                        'Only "int", "long", "string", or "bytes" are allowed'
                        .format(record[Avro.Name], self.current_schema_name, record[Avro.Type])
                    )
