#
# Copyright The NOMAD Authors.
#
# This file is part of NOMAD. See https://nomad-lab.eu for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import io
import typing
import sys

from nomad import config, utils, datamodel

from .api import Auth


def parse(
        mainfile_path: str,
        parser_name: str = None,
        backend_factory: typing.Callable = None,
        strict: bool = True, logger=None):
    '''
    Run the given parser on the downloaded entry. If no parser is given,
    do parser matching and use the respective parser.
    '''
    from nomad import parsing
    from nomad.parsing import parsers
    mainfile = os.path.basename(mainfile_path)

    if logger is None:
        logger = utils.get_logger(__name__)
    if parser_name is not None:
        parser = parsers.parser_dict.get(parser_name)
        assert parser is not None, 'the given parser must exist'
    else:
        mainfile_path = os.path.abspath(mainfile_path)
        parser = parsers.match_parser(mainfile_path, strict=strict)
        if isinstance(parser, parsing.MatchingParser):
            parser_name = parser.name
        else:
            parser_name = parser.__class__.__name__

    assert parser is not None, 'there is no parser matching %s' % mainfile
    logger = logger.bind(parser=parser.name)  # type: ignore
    logger.info('identified parser')
    if hasattr(parser, 'backend_factory'):
        setattr(parser, 'backend_factory', backend_factory)

    entry_archive = datamodel.EntryArchive()
    metadata = entry_archive.m_create(datamodel.EntryMetadata)
    cwd = os.getcwd()
    try:
        mainfile_path = os.path.abspath(mainfile_path)
        os.chdir(os.path.abspath(os.path.dirname(mainfile_path)))
        parser.parse(mainfile_path, entry_archive, logger=logger)
    except Exception as e:
        logger.error('parsing was not successful', exc_info=e)
        raise e
    finally:
        os.chdir(cwd)

    if metadata.domain is None:
        metadata.domain = parser.domain

    logger.info('ran parser')
    return entry_archive


def normalize(
        normalizer: typing.Union[str, typing.Callable], entry_archive, logger=None):
    from nomad import normalizing

    if logger is None:
        logger = utils.get_logger(__name__)

    if isinstance(normalizer, str):
        normalizer = next(
            normalizer_instance for normalizer_instance in normalizing.normalizers
            if normalizer_instance.__class__.__name__ == normalizer)

    assert normalizer is not None, 'there is no normalizer %s' % str(normalizer)
    normalizer_instance = typing.cast(typing.Callable, normalizer)(entry_archive)
    logger = logger.bind(normalizer=normalizer_instance.__class__.__name__)
    logger.info('identified normalizer')

    normalizer_instance.normalize(logger=logger)
    logger.info('ran normalizer')


def normalize_all(entry_archive, logger=None):
    '''
    Parse the downloaded entry and run the whole normalizer chain.
    '''
    from nomad import normalizing

    for normalizer in normalizing.normalizers:
        if normalizer.domain is None or normalizer.domain == entry_archive.metadata.domain:
            normalize(normalizer, entry_archive, logger=logger)


class LocalEntryProcessing:
    '''
    Instances represent a local reproduction of the processing for a single entry.
    It allows to download raw data from a nomad server and reproduce its processing
    (parsing, normalizing) with the locally installed parsers and normalizers.

    The use-case is error/warning reproduction. Use ELK to identify errors, use
    the upload, archive ids to given by ELK, and reproduce and fix the error
    in your development environment.

    Arguments:
        entry_id: The entry_id of the entry to locally process.
        override: Set to true to override any existing local entry data.
        auth: Optional Auth object to download private data.
    '''
    def __init__(self, entry_id: str, override: bool = False, auth: Auth = None) -> None:
        from nomad import files
        from nomad.client import api

        self.entry_id = entry_id
        response = self.__handle_response(
            api.get(f'entries/{self.entry_id}', auth=auth))
        self.mainfile = response.json()['data']['mainfile']
        self.upload_id = response.json()['data']['upload_id']

        self.local_path = os.path.join(config.fs.tmp, f'repro_{self.entry_id}.zip')
        if not os.path.exists(os.path.dirname(self.local_path)):
            os.makedirs(os.path.dirname(self.local_path))

        if not os.path.exists(self.local_path) or override:
            # download raw if not already downloaded or if override is set
            print('Downloading', self.entry_id)
            response = self.__handle_response(
                api.get(f'entries/{self.entry_id}/raw', auth=auth))

            with open(self.local_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=io.DEFAULT_BUFFER_SIZE):
                    f.write(chunk)
        else:
            print('Entry already downloaded.')

        self.upload_files = files.StagingUploadFiles(upload_id=f'tmp_{self.entry_id}', create=True)

    def __handle_response(self, response):
        if response.status_code == 404:
            print('Could not find entry, maybe it was deleted or the ids are wrong:', file=sys.stderr)
            print(response.text, file=sys.stderr)
            sys.exit(1)
        elif response.status_code != 200:
            print(f'API error: {response.text}', file=sys.stderr)
            sys.exit(1)

        return response

    def __enter__(self):
        # open/extract upload file
        print('Extracting entry data.')
        if self.upload_files.is_empty():  # Only add the files once
            self.upload_files.add_rawfiles(self.local_path)

        for raw_file in self.upload_files.raw_directory_list(recursive=True):
            if raw_file.path == f'{self.upload_id}/{self.mainfile}':
                print('Identified mainfile.')
                return self

        print('Count not find file in download.', file=sys.stderr)
        sys.exit(1)

    def __exit__(self, exception_type, exception, trace):
        if exception:
            import traceback
            traceback.print_exc()
        try:
            self.upload_files.delete()
        except Exception:
            pass
        if exception:
            sys.exit(1)

    def parse(self, parser_name: str = None, **kwargs):
        '''
        Run the given parser on the downloaded entry. If no parser is given,
        do parser matching and use the respective parser.
        '''
        return parse(
            self.upload_files.raw_file_object(f'{self.upload_id}/{self.mainfile}').os_path,
            parser_name=parser_name, logger=utils.get_logger(__name__), **kwargs)

    def normalize(self, normalizer: typing.Union[str, typing.Callable], entry_archive=None):
        '''
        Parse the downloaded entry and run the given normalizer.
        '''
        if entry_archive is None:
            entry_archive = self.parse()

        return normalize(
            entry_archive=entry_archive, normalizer=normalizer, logger=utils.get_logger(__name__))

    def normalize_all(self, entry_archive=None):
        '''
        Parse the downloaded entry and run the whole normalizer chain.
        '''
        return normalize_all(
            entry_archive=entry_archive, logger=utils.get_logger(__name__))
