from json import JSONEncoder
from elasticsearch import Elasticsearch
from elasticsearch_dsl import *


def create_typology_filter(value):
    return Q('match', typology=value)


def create_pid_type_filter(value):
    args = {'localIdentifier.type': value}
    return Q('nested', path='localIdentifier', query=Q('bool', must=[Q('match', **args)]))


def create_datasource_filter(value):
    args = {'datasources.datasourceName': value}
    return Q('nested', path='datasources', query=Q('bool', must=[Q('match', **args)]))


class DLIESResponseEncoder(JSONEncoder):
    def default(self, o):
        return o.__dict__


class DLIESResponse(object):
    def __init__(self, facet=None, total=0, hits=[]):
        if facet is None:
            facet = dict(pid=[], typology=[], datasource=[])
        self.facet = facet
        self.total = total
        self.hits = hits


class DLIESConnector(object):
    def __init__(self, index_host, index_name):
        self.index_host = index_host
        self.client = Elasticsearch(hosts=[index_host])
        self.index_name = index_name

    def simple_query(self, textual_query, start=None, end=None, filter_key=None, filter_value=None):
        s = Search(using=self.client, index=self.index_name).doc_type('object').sort('-relatedDatasets',
                                                                                     '-relatedPublications',
                                                                                     '-relatedUnknown')
        q = Q('match', _all=textual_query)
        s.aggs.bucket('typologies', 'terms', field='typology')
        s.aggs.bucket('all_datasources', 'nested', path='datasources').bucket('all_names', 'terms',
                                                                              field='datasources.datasourceName')

        s.aggs.bucket('all_pids', 'nested', path='localIdentifier').bucket('all_types', 'terms',
                                                                           field='localIdentifier.type')

        if filter_key is not None and len(filter_key) > 0:
            if filter_key == 'typology':
                s = s.query(q).filter(create_typology_filter('dataset'))
            elif filter_key == 'datasource':
                s = s.query(q).filter(create_datasource_filter(filter_value))
            elif filter_key == 'pidtype':
                s = s.query(q).filter(create_pid_type_filter(filter_value))
            else:
                s = s.query(q)
        else:
            s = s.query(q)

        if start is not None:
            if end is None:
                end = start + 10
            s = s[start:end]
        response = s.execute()

        hits = []

        for index_result in response.hits:
            hits.append(index_result.__dict__['_d_'])

        pid_types = []
        for tag in response.aggs.all_pids.all_types.buckets:
            pid_types.append(dict(key=tag.key, count=tag.doc_count))

        datasources = []
        for tag in response.aggs.all_datasources.all_names.buckets:
            datasources.append(dict(key=tag.key, count=tag.doc_count))

        typologies = []
        for tag in response.aggs.typologies.buckets:
            typologies.append(dict(key=tag.key, count=tag.doc_count))

        return DLIESResponse(total=response.hits.total,
                             facet=dict(pid=pid_types, typology=typologies, datasource=datasources), hits=hits)

    def related_type(self, object_id, object_type, start=None):
        args = {'target.objectType': object_type}
        query_type = Q('nested', path='target', query=Q('bool', must=[Q('match', **args)]))
        args_id = {'source.dnetIdentifier': object_id}
        query_for_id = Q('nested', path='source', query=Q('bool', must=[Q('match', **args_id)]))
        s = Search(using=self.client).doc_type('scholix').query(query_for_id & query_type)
        if start:
            s = s[start:start + 10]

        response = s.execute()
        hits = []

        for index_hit in response.hits:
            hits.append(index_hit.__dict__['_d_'])

        return hits

    def fix_collectedFrom(self, source, relation):
        relSource = relation.get('source')
        collectedFrom = relSource['collectedFrom']
        for coll in collectedFrom:
            for d in source['datasources']:
                if d['datasourceName'] == coll['provider']['name']:
                    d['provisionMode'] = coll['provisionMode']
        return source

    def item_by_id(self, id, type=None, start=None):
        try:
            res = self.client.get(index=self.index_name, doc_type='object', id=id)
            hits = []
            input_source = res['_source']
            related_publications = []
            related_dataset = []
            related_unknown = []

            rel_source = None
            if input_source.get('relatedPublications') > 0:
                if 'publication' == type:
                    related_publications = self.related_type(id, 'publication', start)
                else:
                    related_publications = self.related_type(id, 'publication')
                rel_source = related_publications[0]
            if input_source.get('relatedDatasets') > 0:
                if 'dataset' == type:
                    related_dataset = self.related_type(id, 'dataset', start)
                else:
                    related_dataset = self.related_type(id, 'dataset')
                rel_source = related_dataset[0]
            if input_source.get('relatedUnknown') > 0:
                if 'unknown' == type:
                    related_unknown = self.related_type(id, 'unknown', start)
                else:
                    related_unknown = self.related_type(id, 'unknown')
                rel_source = related_unknown[0]

            input_source = self.fix_collectedFrom(input_source, rel_source)
            hits.append(input_source)

            hits.append(dict(related_publications=related_publications, related_dataset=related_dataset,
                             related_unknown=related_unknown))

            return DLIESResponse(total=1, hits=hits)
        except:
            return DLIESResponse()

# connector = DLIESConnector('node0-d-dli.d4science.org', 'dli')
# res = connector.item_by_id('70|r3d100010134::819365f759cbc38c0362fe8bc684ee7b','dataset',100)
# for hit in res.hits[1]['related_dataset']:
#     print hit['dnetIdTarget']


# res = connector.item_by_id('70|r3d100010134::819365f759cbc38c0362fe8bc684ee7b','dataset',1000)
# for hit in res.hits[1]['related_dataset']:
#     print hit['dnetIdTarget']
# print json.dumps(res, cls=DLIESResponseEncoder, indent=2)