from django import forms

from haystack.forms import SearchForm

class IconolabSearchForm(SearchForm):

	model_type = forms.ChoiceField(required=False, choices=(("images","Image"), ("annotations","Annotation")) )

	def __init__(self, *args, **kwargs):

		self.collection_name = kwargs.pop("collection_name")

		selected_model_type =  kwargs.pop("model_type", None)

		if selected_model_type is not None: 

			data = kwargs.get("data", None)

			if data:

				data = data.copy()

				data["model_type"] = selected_model_type

				kwargs['data'] = data

		super(IconolabSearchForm, self).__init__(*args, **kwargs)

	def no_query_found(self):

		# load all

		selected_type = self.cleaned_data.get("model_type")

		qs = self.get_model_type_queryset(self.searchqueryset, selected_type).load_all()

		return qs

	def get_model_type_queryset(self, qs, model_type):

		if model_type == 'images':

		if model_type == 'annotations':

		if self.collection_name is not None:

			qs = qs.filter(collection = self.collection_name)

		return qs

	def search(self):

		selected_type = self.cleaned_data.get("model_type")

		qs = super(IconolabSearchForm, self).search()

		if qs.count() == 0:

			return qs

		else:

			qs = self.get_model_type_queryset(qs, selected_type).load_all()

		return qs