"""
Container for Avro schemas.
"""

import copy
import json

from collections import OrderedDict
from json import JSONDecodeError
from pathlib import Path
from typing import Optional, Union, cast
from ruamel.yaml import YAML  # type: ignore

from avro_preprocessor.avro_domain import Avro
from avro_preprocessor.avro_naming import AvroNaming
from avro_preprocessor.avro_paths import AvroPaths

from avro_preprocessor.colored_json import ColoredJson

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


class SchemasContainer:
    """
    A memory container from schemas.
    """

    def __init__(
            self,
            paths: AvroPaths,
            verbose: bool = True,
    ) -> None:
        self.paths: AvroPaths = paths
        self.verbose: bool = verbose
        self.original: OrderedDict = OrderedDict()
        self.processed: OrderedDict = OrderedDict()
        self.autogenerated_keys: OrderedDict = OrderedDict()

        self.modules: OrderedDict = OrderedDict()

        self.original_string: OrderedDict = OrderedDict()

    def read_schemas(self) -> None:
        """
        Reads schemas from disk
        """
        self.paths.traverse_path(
            self.paths.input_path, self.paths.input_schema_file_extension, self._get_schema)

        self.processed = copy.deepcopy(self.original)

        if self.verbose:
            print('')

    def _get_schema(self, path: Path) -> None:
        schema_text = path.read_text()
        try:
            try:
                schema = json.loads(schema_text, object_pairs_hook=OrderedDict)
            except JSONDecodeError:
                yaml = YAML()
                schema = yaml.load(schema_text)
                # print(ColoredJson.highlight_yaml(schema))

        except Exception as e:
            import sys
            raise type(e)(str(e) + '\nError while processing schema %s' % path) \
                .with_traceback(sys.exc_info()[2])
        namespace = str(
            Path(
                str(path.absolute())
                .replace(str(Path(self.paths.input_path).resolve()) + '/', '')
                .replace('/', '.')
            )
            .with_suffix('')
        )

        self._store_schema(schema, namespace)
        self.original_string[namespace] = schema_text

        if self.verbose:
            print('Reading', path, ' -> ', namespace)

    def _store_schema(self, schema: OrderedDict, namespace: str) -> None:
        schema_doc = schema[Avro.Doc] if isinstance(schema[Avro.Doc], str) \
            else ' '.join(schema[Avro.Doc])

        # if we run the preprocessor with the KeyGenerator module and then re-run it again
        # we need to recognize the autogenerated key schemas and store them
        # in the right container
        if schema_doc.startswith(AvroNaming.autogenerated_key_doc):
            self.autogenerated_keys[namespace] = schema
        else:
            self.original[namespace] = schema

    def write_schemas(self) -> None:
        """
        Writes schemas to disk.
        """
        # let's cleanup the output schema directory first
        AvroPaths.reset_directory(self.paths.output_path)

        schemas_and_keys = list(self.processed.items()) + list(self.autogenerated_keys.items())

        for name, schema in schemas_and_keys:
            path = self.paths.to_output_path(name)
            path.parent.mkdir(parents=True, exist_ok=True)

            if self.paths.output_schema_file_format.lower() == 'yaml':
                yaml = YAML()
                yaml.default_flow_style = False
                yaml.indent(**ColoredJson.yaml_indent)
                yaml.Representer.add_representer(OrderedDict, yaml.Representer.represent_dict)
                yaml.dump(schema, path)
            else:
                schema_text = json.dumps(schema, indent=ColoredJson.json_indent)
                path.write_text(schema_text)

    def print_schemas(self, how: str = 'json', indent: Union[None, int, dict] = None) -> None:
        """
        Print all schemas.
        :param how: 'json' or 'yaml' or 'dict'
        :param indent: how much to indent printed JSONs or YAMLs
        """

        def print_schema(schema_name: str, data: OrderedDict) -> None:
            value: Union[str, OrderedDict]
            if how == 'json':
                value = ColoredJson.highlight_json(data, cast(Optional[int], indent))
            elif how == 'yaml':
                value = ColoredJson.highlight_yaml(data, cast(Optional[dict], indent))
            else:  # 'dict'
                value = data
            print(schema_name, value)

        print('#### ORIGINAL SCHEMAS ####')
        for name, schema in self.original.items():
            print_schema(name, schema)

        print('#### PROCESSED SCHEMAS ####')
        for name, schema in self.processed.items():
            print_schema(name, schema)
