#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Usage:  program.py <customer>

"""

from typing import Mapping, Any, List, MutableSequence
import copy
import json
import sys
import os

import bel.lang.belobj
import bel.lang.bel_specification
import bel.db.arangodb as arangodb
import bel.utils as utils


from bel.Config import config

import structlog
log = structlog.getLogger(__name__)

Edges = MutableSequence[Mapping[str, Any]]

arango_client = arangodb.get_client()
edgestore_db = arangodb.get_edgestore_handle(arango_client)


def nanopub_to_edges(nanopub: dict = {}, rules: List[str] = [], orthologize_targets: list = []):
    """Process nanopub into edges and load into EdgeStore

    Args:
        nanopub: BEL Nanopub
        rules: list of compute rules to process
        orthologize_targets: list of species in TAX:<int> format

    Returns:
        list: of edges

    Edge object:
        {
            "edge": {
                "subject": {
                    "name": subj_canon,
                    "name_lc": subj_canon.lower(),
                    "label": subj_lbl,
                    "label_lc": subj_lbl.lower(),
                    "components": subj_components,
                },
                "relation": {  # relation _key is based on a hash
                    "relation": edge_ast.bel_relation,
                    "edge_hash": edge_hash,
                    "edge_dt": edge_dt,
                    "nanopub_url": nanopub_url,
                    "nanopub_id": nanopub_id,
                    "citation": citation,
                    "subject_canon": subj_canon,
                    "subject": subj_lbl,
                    "object_canon": obj_canon,
                    "object": obj_lbl,
                    "annotations": nanopub['annotations'],
                    "metadata": nanopub['metadata'],
                    "public_flag": True,  # will be added when groups/permissions feature is finished,
                    "edge_types": edge_types,
                },
                'object': {
                    "name": obj_canon,
                    "name_lc": obj_canon.lower(),
                    "label": obj_lbl,
                    "label_lc": obj_lbl.lower(),
                    "components": obj_components,
                }
            }
        }
    """

    # Collect input values ####################################################
    nanopub_url = nanopub.get('source_url', '')

    edge_dt = utils.dt_utc_formatted()  # don't want this in relation_id

    # Extract BEL Version and make sure we can process this
    if nanopub['nanopub']['type']['name'].upper() == "BEL":
        bel_version = nanopub['nanopub']['type']['version']
        versions = bel.lang.bel_specification.get_bel_versions()
        if bel_version not in versions:
            log.error(f'Do not know this BEL Version: {bel_version}, these are the ones I can process: {versions.keys()}')
            return []
    else:
        log.error(f"Not a BEL Nanopub according to nanopub.type.name: {nanopub['nanopub']['type']['name']}")
        return []

    # Required for BEL parsing/canonicalization/orthologization
    api_url = config['bel_api']['servers']['api_url']

    try:
        citation_string = normalize_nanopub_citation(nanopub)
    except Exception as e:
        log.error(f'Could not create citation string for {nanopub_url}')
        citation_string = ''

    if orthologize_targets == []:
        if config['bel_api'].get('edges', None):
            orthologize_targets = config['bel_api']['edges'].get('orthologize_targets', [])

    # orig_species_id = [anno['id'] for anno in nanopub['nanopub']['annotations'] if anno['type'] == 'Species']
    # if orig_species_id:
    #     orig_species_id = orig_species_id[0]

    master_annotations = copy.deepcopy(nanopub['nanopub']['annotations'])
    master_metadata = copy.deepcopy(nanopub['nanopub']['metadata'])
    master_metadata.pop('gd:abstract', None)

    nanopub_type = nanopub['nanopub']['metadata'].get('nanopub_type')

    # Create Edge Assertion Info ##############################################
    # r = generate_assertion_edge_info(nanopub['nanopub']['assertions'], orig_species_id, orthologize_targets, bel_version, api_url, nanopub_type)
    r = generate_assertion_edge_info(nanopub['nanopub']['assertions'], orthologize_targets, bel_version, api_url, nanopub_type)
    edge_info_list = r['edge_info_list']

    # Build Edges #############################################################
    edges = []
    errors = []
    for edge_info in edge_info_list:
        annotations = copy.deepcopy(master_annotations)
        metadata = copy.deepcopy(master_metadata)

        errors.extend(edge_info['errors'])

        if not edge_info.get('canonical'):
            continue

        # TODO - remove this
        # if edge_info.get('species_id', False):
        #     annotations = orthologize_context(edge_info['species_id'], annotations)

        edge_hash = utils._create_hash(f'{edge_info["canonical"]["subject"]} {edge_info["canonical"]["relation"]} {edge_info["canonical"]["object"]}')

        edge = {
            'edge': {
                'subject': {
                    "name": edge_info['canonical']['subject'],
                    "name_lc": edge_info['canonical']['subject'].lower(),
                    "label": edge_info['decanonical']['subject'],
                    "label_lc": edge_info['decanonical']['subject'].lower(),
                    "components": edge_info['subject_comp'],
                },
                'relation': {
                    'relation': edge_info['canonical']['relation'],
                    'edge_hash': edge_hash,
                    'edge_dt': edge_dt,
                    'nanopub_url': nanopub_url,
                    'nanopub_id': nanopub['nanopub']['id'],
                    'citation': citation_string,
                    'subject_canon': edge_info['canonical']['subject'],
                    'subject': edge_info['decanonical']['subject'],
                    'object_canon': edge_info['canonical']['object'],
                    'object': edge_info['decanonical']['object'],
                    'annotations': copy.deepcopy(annotations),
                    'metadata': copy.deepcopy(metadata),
                    'public_flag': True,
                    'edge_types': edge_info['edge_types'],
                    'species_id': edge_info['species_id'],
                    'species_label': edge_info['species_label'],
                },
                'object': {
                    "name": edge_info['canonical']['object'],
                    "name_lc": edge_info['canonical']['object'].lower(),
                    "label": edge_info['decanonical']['object'],
                    "label_lc": edge_info['decanonical']['object'].lower(),
                    "components": edge_info['object_comp'],
                }
            }
        }

        edges.append(copy.deepcopy(edge))

    return {"edges": edges, "nanopub_id": nanopub['nanopub']['id'], "nanopub_url": nanopub_url, "success": True, "errors": errors}


def extract_ast_species(ast):
    """Extract species from ast.species set of tuples (id, label)"""

    species_id = 'None'
    species_label = 'None'
    species = [(species_id, species_label,) for (species_id, species_label) in ast.species if species_id]

    if len(species) == 1:
        (species_id, species_label) = species[0]

    if not species_id:
        species_id = 'None'
        species_label = 'None'

    log.debug(f'AST Species: {ast.species}  Species: {species}  SpeciesID: {species_id}')

    return (species_id, species_label)


# def generate_assertion_edge_info(assertions: List[dict], orig_species_id: str, orthologize_targets: List[str], bel_version: str, api_url: str, nanopub_type: str) -> dict:
def generate_assertion_edge_info(assertions: List[dict], orthologize_targets: List[str], bel_version: str, api_url: str, nanopub_type: str = '') -> dict:
    """Create edges (SRO) for assertions given orthologization targets

    Args:
        assertions: list of BEL statements (SRO object)
        orthologize_targets: list of species in TAX:<int> format
        bel_version: to be used for processing assertions
        api_url: BEL API url endpoint to use for terminologies and orthologies
    """

    bo = bel.lang.belobj.BEL(bel_version, api_url)
    bo_computed = bel.lang.belobj.BEL(bel_version, api_url)

    edge_info_list = []

    with utils.Timer() as t:
        for assertion in assertions:
            # if not assertion.get('relation', False):
            #     continue  # Skip any subject only statements

            start_time = t.elapsed
            bo.parse(assertion)

            if not bo.ast:
                errors = [f'{error[0]} {error[1]}' for error in bo.validation_messages if error[0] == 'ERROR']
                edge_info = {'errors': copy.deepcopy(errors)}
                edge_info_list.append(copy.deepcopy(edge_info))
                continue

            # populate canonical terms and orthologs for assertion
            bo.collect_nsarg_norms()
            bo.collect_orthologs(orthologize_targets)

            log.info('Timing - time to collect nsargs and orthologs', delta_ms=(t.elapsed - start_time))

            (edge_species_id, edge_species_label) = extract_ast_species(bo.ast)
            orig_species_id = edge_species_id

            canon = bo.canonicalize().to_triple()

            components = get_node_subcomponents(bo.ast)  # needs to be run after canonicalization
            computed_asts = bo.compute_edges(ast_result=True)  # needs to be run after canonicalization
            decanon = bo.decanonicalize().to_triple()

            if nanopub_type == 'backbone':
                edge_types = ['backbone']
            else:
                edge_types = ['original', 'primary']

            if assertion.get('relation', False):

                causal_edge_type = []
                if 'causal' in bo.spec['relations']['info'][assertion['relation']]['categories']:
                    edge_types.append('causal')

                edge_info = {
                    'edge_types': edge_types,
                    'species_id': edge_species_id,
                    'species_label': edge_species_label,
                    'canonical': canon,
                    'decanonical': decanon,
                    'subject_comp': components['subject_comp'],
                    'object_comp': components['object_comp'],
                    'errors': [],
                }
                edge_info_list.append(copy.deepcopy(edge_info))

            # Loop through primary computed asts
            for computed_ast in computed_asts:
                bo_computed.ast = computed_ast
                bo_computed.collect_nsarg_norms()
                canon = bo_computed.canonicalize().to_triple()
                components = get_node_subcomponents(bo_computed.ast)  # needs to be run after canonicalization
                decanon = bo_computed.decanonicalize().to_triple()

                (edge_species_id, edge_species_label) = extract_ast_species(bo_computed.ast)

                edge_info = {
                    'edge_types': ['computed'],
                    'species_id': edge_species_id,
                    'species_label': edge_species_label,
                    'canonical': canon,
                    'decanonical': decanon,
                    'subject_comp': components['subject_comp'],
                    'object_comp': components['object_comp'],
                    'errors': [],
                }
                if [edge for edge in edge_info_list if edge.get('canonical', {}) == canon]:
                    continue  # skip if edge is already included (i.e. the primary is same as computed edge)
                edge_info_list.append(copy.deepcopy(edge_info))

            # Skip orthologs if backbone nanopub
            if nanopub_type == 'backbone':
                continue

            # only process orthologs if there are species-specific NSArgs
            if len(bo.ast.species) > 0:
                # Loop through orthologs
                for species_id in orthologize_targets:
                    log.debug(f'Orig species: {orig_species_id}  Target species: {species_id}')
                    if species_id == orig_species_id:
                        continue

                    bo.orthologize(species_id)

                    (edge_species_id, edge_species_label) = extract_ast_species(bo.ast)

                    if edge_species_id == 'None' or edge_species_id == orig_species_id:
                        log.debug(f'Skipping orthologization- species == "None" or {orig_species_id}  ASTspecies: {bo.ast.species} for {bo}')
                        continue

                    ortho_decanon = bo.orthologize(species_id).to_triple()  # defaults to decanonicalized orthologized form
                    ortho_canon = bo.canonicalize().to_triple()
                    computed_asts = bo.compute_edges(ast_result=True)  # needs to be run after canonicalization
                    components = get_node_subcomponents(bo.ast)  # needs to be run after canonicalization

                    if assertion.get('relation', False):
                        edge_info = {
                            'edge_types': ['orthologized', 'primary'] + causal_edge_type,
                            'species_id': edge_species_id,
                            'species_label': edge_species_label,
                            'canonical': ortho_canon,
                            'decanonical': ortho_decanon,
                            'subject_comp': components['subject_comp'],
                            'object_comp': components['object_comp'],
                            'errors': [],
                        }

                        edge_info_list.append(copy.deepcopy(edge_info))

                    # Loop through orthologized computed asts
                    for computed_ast in computed_asts:
                        bo_computed.ast = computed_ast
                        bo_computed.collect_nsarg_norms()
                        canon = bo_computed.canonicalize().to_triple()
                        components = get_node_subcomponents(bo_computed.ast)  # needs to be run after canonicalization
                        decanon = bo_computed.decanonicalize().to_triple()

                        (edge_species_id, edge_species_label) = extract_ast_species(bo_computed.ast)

                        edge_info = {
                            'edge_types': ['computed', 'orthologized'],
                            'species_id': edge_species_id,
                            'species_label': edge_species_label,
                            'canonical': canon,
                            'decanonical': decanon,
                            'subject_comp': components['subject_comp'],
                            'object_comp': components['object_comp'],
                            'errors': [],
                        }
                        if [edge for edge in edge_info_list if edge.get('canonical', {}) == canon]:
                            continue  # skip if edge is already included (i.e. the primary is same as computed edge)
                        edge_info_list.append(copy.deepcopy(edge_info))

    log.info('Timing - Generated all edge info for all nanopub assertions', delta_ms=t.elapsed)

    return {'edge_info_list': edge_info_list}


def get_node_subcomponents(ast):

    sub, obj = [], []

    try:
        sub = ast.bel_subject.subcomponents(subcomponents=[])

        # TODO - update handling nested BEL statement - see Natalie's recommendation on how to handle
        obj_components = []
        if ast.bel_object.__class__.__name__ == 'BELAst':  # Nested BEL Assertion
            obj_components = ast.bel_object.bel_subject.subcomponents(subcomponents=[])
            obj_components = ast.bel_object.bel_object.subcomponents(subcomponents=obj_components)
        elif hasattr(ast.bel_object, 'subcomponents'):  # Normal BEL Assertion
            obj_components = ast.bel_object.subcomponents(subcomponents=[])

        obj = obj_components

    except Exception as e:
        log.warning(f'Problem getting subcomponents for {ast.to_string()}', error=str(e))

    return {'subject_comp': sub, 'object_comp': obj}


def orthologize_context(orthologize_target: str, annotations: Mapping[str, Any]) -> Mapping[str, Any]:
    """Orthologize context

    Replace Species context with new orthologize target and add a annotation type of OrthologizedFrom
    """

    url = f'{config["bel_api"]["servers"]["api_url"]}/terms/{orthologize_target}'
    r = utils.get_url(url)
    species_label = r.json().get("label", "unlabeled")

    orthologized_from = {}
    for idx, annotation in enumerate(annotations):
        if annotation['type'] == 'Species':
            orthologized_from = {'id': annotation['id'], 'label': annotation['label']}
            annotations[idx] = {'type': 'Species', 'id': orthologize_target, 'label': species_label}

    if 'id' in orthologized_from:
        annotations.append({'type': 'OrigSpecies', 'id': f'Orig-{orthologized_from["id"]}', 'label': f'Orig-{orthologized_from["label"]}'})

    return annotations


def normalize_nanopub_citation(nanopub):

    citation_string = ''
    if 'database' in nanopub['nanopub']['citation']:
        if nanopub['nanopub']['citation']['database']['name'].lower() == 'pubmed':
            citation_string = f"PMID:{nanopub['nanopub']['citation']['database']['id']}"
        else:
            citation_string = f"{nanopub['nanopub']['citation']['database']['name']}:{nanopub['nanopub']['citation']['database']['id']}"
    elif 'reference' in nanopub['nanopub']['citation'] and isinstance(nanopub['nanopub']['citation']['reference'], str):
        citation_string = nanopub['nanopub']['citation']['reference']
    elif 'uri' in nanopub['nanopub']['citation']:
        citation_string = nanopub['nanopub']['citation']['uri']

    return citation_string


def main():
    pass


if __name__ == '__main__':
    # Setup logging
    import logging.config
    module_fn = os.path.basename(__file__)
    module_fn = module_fn.replace('.py', '')

    if config.get('logging', False):
        logging.config.dictConfig(config.get('logging'))

    log = logging.getLogger(f'{module_fn}')

    main()
