from datetime import datetime
import trafaret as t

from django.core.exceptions import MultipleObjectsReturned, ObjectDoesNotExist

from staff.person.models import Staff

from .constants import OWNER
from .models import ContractorFirm, Badge


def file_validator():
    return t.Any()


def ChoiceValidator(choice):
    return t.Or(*(t.Atom(x[0]) for x in choice))


CodeValidator = t.String(regex='[0-9]+')


def select_model(manager, key, related=None):
    def wrapped(val):
        try:
            if related:
                return manager.select_related(*related).get(**{key: val})
            return manager.get(**{key: val})
        except ObjectDoesNotExist:
            raise t.DataError('Not found')
        except MultipleObjectsReturned:
            # This is not ok, lets 500
            raise
    return wrapped


def check_ownertype(otype):
    def wrapped(badge):
        if badge.ownertype != otype:
            raise t.DataError('Wrong owner type: not %s' % otype)
        return badge
    return wrapped


StaffLoginValidator = t.String() >> select_model(
    Staff.objects, 'login', related=('office__city__country', ))

OrgValidator = t.Int(gt=0) >> select_model(ContractorFirm.objects, 'id')

CandidateValidator = (
    t.Int(gt=0) >> select_model(Badge.objects.filter(owner=OWNER.CANDIDATE), 'id'))


def extract_date(s, format):
    try:
        return datetime.strptime(s, format).date()
    except ValueError as exc:
        raise t.DataError(exc)
