import json
import queue
import logging
import traceback
import multiprocessing
from heapq import nsmallest
from collections import Counter
from collections import OrderedDict
from collections import defaultdict
from itertools import chain
from itertools import zip_longest
from itertools import cycle
from operator import itemgetter

from . import score
from . import storage
from . import featurizers


first = itemgetter(0)
logger = logging.getLogger(__name__)
cat = chain.from_iterable


class defaults:
    n = 600
    r = int(1e5)
    k = None
    limit = 5


class Index(object):
    def __init__(self, backend):
        self.backend = backend
        self.rowcount = self.backend.get_rowcount()
        self.featurizer = featurizers.all[self.backend.featurizer_name]

    def _search(self, query, r, n, k):
        toks = self.featurizer(query)
        toks = self.backend.find_least_frequent_tokens(toks, r, k)
        r_map = Counter()
        for i, tok in enumerate(toks, 1):
            rng = self.backend.get_token(tok)
            r_map.update(rng)
        top_ids = map(first, r_map.most_common(n))
        return list(top_ids)

    def _scored_records(self, record_ids, orig_query,
                        extract_func=score.features, score_func=score.hit):
        orig_features = extract_func(orig_query)
        for rownum, r in zip(record_ids, self.backend.get_records(record_ids)):
            s = score_func(orig_features, extract_func(r.fields))
            yield s, rownum, r

    def search(self, query, limit=defaults.limit, r=defaults.r, n=defaults.n,
               k=None, extract_func=score.features, score_func=score.hit):
        record_ids = self._search(query, r, n, k)
        scores_records = self._scored_records(
            record_ids, query, extract_func, score_func)
        return [
            {"fields": rec.fields, "pk": rec.pk, "score": s,
             "data": rec.data, "rownum": rownum}
            for s, rownum, rec in nsmallest(limit, scores_records, key=first)
        ]

    def _save_records(self, records, idxs=[]):
        completed = []
        for rec, i in zip_longest(records, idxs):
            try:
                idx = (self.backend.save_record(rec) if i is None
                       else self.backend.update_record(rec, i))
            except:
                for idx in completed:
                    self.backend.delete_record(idx)
                raise
            completed.append(idx)
        self.backend.increment_rowcount(len(completed))
        return completed

    def _update_tokens(self, tokmap, freq_update):
        for tok in tokmap.keys():
            idxs = tokmap[tok]
            self.backend.update_token(tok, idxs)
            freq_update[tok] = len(idxs)

    def _update_tokens_and_freqs(self, tokmap):
        freq_update = {}
        try:
            self._update_tokens(tokmap, freq_update)
            self.backend.update_freqs(freq_update.items())
        except:
            for tok in freq_update:
                self.backend.drop_records_from_token(tok, tokmap[tok])
            raise

    def add(self, records, idxs=[]):
        idxs = list(self._save_records(records, idxs))
        tokmap = defaultdict(list)
        for idx, rec in zip(idxs, records):
            for tok in self.featurizer(rec.fields):
                tokmap[tok].append(idx)
        self._update_tokens_and_freqs(tokmap)
        return idxs

    def close(self):
        return self.backend.close()


class ParallelIndexWorker(multiprocessing.Process):
    def __init__(self, backend_name, work_q, result_q):
        super().__init__()
        self.work_q = work_q
        self.result_q = result_q
        self.counters = dict()
        self.be_cls = storage.backends[backend_name]
        self.methods = dict(count_tokens=self._count_tokens,
                            score_records=self._score_records)

    def _count_tokens(self, query_key, n_total_toks, blob, n):
        r_map = self.counters.get(query_key, None)
        if not r_map:
            r_map = self.counters[query_key] = {'cntr': Counter(),
                                                'n_toks': 0}
        if blob:
            r_map['cntr'].update(self.be_cls._get_token(blob))
        r_map['n_toks'] += 1
        if r_map['n_toks'] >= n_total_toks:
            cnt = self.counters.pop(query_key)['cntr']
            return list(map(first, cnt.most_common(n)))
        else:
            return None

    def _scores(self, blobs, orig_features,
                extract_func=score.features, score_func=score.hit):
        for rownum, blob in blobs:
            r = self.be_cls._get_record(blob)
            s = score_func(orig_features, extract_func(r.fields))
            yield s, rownum, r

    def _score_records(self, orig_features, limit, blobs,
                       extract_func=score.features, score_func=score.hit):
        scores_records = self._scores(blobs, orig_features,
                                      extract_func, score_func)
        return nsmallest(limit, scores_records, key=first)

    def run(self):
        while True:
            try:
                logger.debug("Getting work")
                query_id, meth, args = self.work_q.get()
            except IOError as e:
                logger.debug("Received IOError (%s) errno %s from work_q",
                             e.message, e.errno)
                break
            except EOFError:
                logger.debug("Received EOFError from work_q")
                break
            if meth == 'stop':
                logger.debug("Received sentinel, stopping")
                break
            try:
                logger.debug("Running method %s", meth)
                ret = self.methods[meth](*args)
            except Exception as e:
                ret = e