from collections import defaultdict

from .db import WikidataEntityLabel, WikidataClaim


class EntityLabelLookup:

    def __init__(self):
        self.cache = {}

    def lookup(self, entity):
        if entity in self.cache:
            return self.cache[entity]
        else:
            item = WikidataEntityLabel.get_or_none(WikidataEntityLabel.entity_id == entity,
                                                   WikidataEntityLabel.language == 'en')
            if item:
                self.cache[entity] = item.value
            else:
                self.cache[entity] = entity
            return self.cache[entity]


label_lookup = EntityLabelLookup()


def get_target_label(t):
    return label_lookup.lookup(t.target_entity)


def get_target_value(claim):
    return claim.body


def get_target_text_value(claim):
    return (
        claim.body
        .get('mainsnak', {})
        .get('datavalue', {})
        .get('value', {})
        .get('text')
    )


def default_agg(x):
    return x


def gather(source_entity, prop_id, map_fn, agg_fn):
    claim = WikidataClaim.select().where(
        WikidataClaim.source_entity == source_entity,
        WikidataClaim.property_id == prop_id,
    )

    if map_fn is None:
        map_fn = get_target_label

    if agg_fn is None:
        agg_fn = default_agg

    return agg_fn([map_fn(c) for c in claim])


def fetch(name, prop_id, map_fn=None, agg_fn=None):
    return (name, prop_id, map_fn, agg_fn)


def gather_all(source_entity, cfg):
    item = {}
    for (name, prop_id, map_fn, agg_fn) in cfg:
        item[name] = gather(source_entity, prop_id, map_fn, agg_fn)
    return item
