from datetime import timedelta, datetime
from six import iteritems, itervalues
from functools import partial

from .search import Search
from .filter import F
from .aggs import Terms, DateHistogram, Histogram
from .utils import AttrDict
from .result import Response

DATE_INTERVALS = {
    'month': lambda d: (d+timedelta(days=32)).replace(day=1),
    'week': lambda d: d+timedelta(days=7),
    'day': lambda d: d+timedelta(days=1),
    'hour': lambda d: d+timedelta(hours=1),

}

AGG_TO_FILTER = {
    Terms: lambda a, v: F('term', **{a.field: v}),
    DateHistogram: lambda a, v: F('range', **{a.field: {'gte': v, 'lt': DATE_INTERVALS[a.interval](v)}}),
    Histogram: lambda a, v:  F('range', **{a.field: {'gte': v, 'lt': v+a.interval}}),
}

BUCKET_TO_DATA = {
    Terms: lambda bucket, filter: (bucket['key'], bucket['doc_count'], bucket['key'] in filter),
    DateHistogram: lambda bucket, filter: (datetime.utcfromtimestamp(int(bucket['key']) / 1000), bucket['doc_count'], bucket['key'] in filter)
}

def agg_to_filter(agg, value):
    return AGG_TO_FILTER[agg.__class__](agg, value)


class FacetedResponse(Response):
    def __init__(self, search, *args, **kwargs):
        super(FacetedResponse, self).__init__(*args, **kwargs)
        super(AttrDict, self).__setattr__('_search', search)

    @property
    def facets(self):
        if not hasattr(self, '_facets'):
            super(AttrDict, self).__setattr__('_facets', AttrDict({}))
            for name, agg in iteritems(self._search.facets):
                buckets = self._facets[name] = []
                data = self.aggregations['_filter_' + name][name]['buckets']
                filter = self._search._raw_filters.get(name, ())
                for b in data:
                    buckets.append(BUCKET_TO_DATA[agg.__class__](b, filter))
        return self._facets



class FacetedSearch(object):
    index = '_all'
    doc_types = ['_all']
    fields = ('*', )
    facets = {}

    def __init__(self, query=None, filters={}):
        self._query = query
        self._raw_filters = {}
        self._filters = {}
        for name, value in iteritems(filters):
            self.add_filter(name, value)

    def add_filter(self, name, value):
        agg = self.facets[name]
        if isinstance(value, list):
            self._filters[name] = F('bool', should=[agg_to_filter(agg, v) for v in value])
            self._raw_filters[name] = value
        else:
            self._filters[name] = agg_to_filter(agg, value)
            self._raw_filters[name] = (value, )

    def search(self):
        return Search(doc_type=self.doc_types, index=self.index)

    def query(self, search):
        if self._query:
            return search.query('multi_match', fields=self.fields, query=self._query)
        return search

    def aggregate(self, search):
        for f, agg in iteritems(self.facets):
            agg_filter = F('match_all')
            for field, filter in iteritems(self._filters):
                if f == field:
                    continue
                agg_filter &= filter
            search.aggs.bucket(
                '_filter_' + f,
                'filter',
                filter=agg_filter
            ).bucket(f, agg)

    def filter(self, search):
        post_filter = F('match_all')
        for f in itervalues(self._filters):
            post_filter &= f
        return search.post_filter(post_filter)

    def highlight(self, search):
        """
        Add highlighting for all the fields
        """
        return search.highlight(*self.fields)

    def build_search(self):
        s = self.search()
        s = self.query(s)
        self.aggregate(s)
        s = self.filter(s)
        s = self.highlight(s)
        return s

    def execute(self):
        if not hasattr(self, '_response'):
            s = self.build_search()
            self._response = s.execute(response_class=partial(FacetedResponse, self))

        return self._response


@@ -2637,32 +2637,39 @@
 ery(self, search
+, query
 ):%0A        if se
@@ -2667,17 +2667,137 @@
     
-if self._
+%22%22%22%0A        Add query part to %60%60search%60%60.%0A%0A        Override this if you wish to customize the query used.%0A        %22%22%22%0A        if 
 quer
@@ -2872,22 +2872,16 @@
 , query=
-self._
 query)%0A 
@@ -3687,32 +3687,97 @@
 d_search(self):%0A
+        %22%22%22%0A        Construct the %60%60Search%60%60 object.%0A        %22%22%22%0A
         s = self
@@ -3814,34 +3814,21 @@
 ry(s
-)%0A        self.aggregate(s
+, self._query
 )%0A  
@@ -3878,24 +3878,50 @@
 ighlight(s)%0A
+        self.aggregate(s)%0A
         retu

Add query parameter for easier overriding