# -*- coding: UTF-8 -*-

from gm_dataquery.dataquery import DataSQLQuery
from gm_dataquery.db import DB

from gm_types.wiki import ITEM_H1, DRUG_H1, MACHINE_H1, MATERIAL_H1, DIRECTORY_TYPE

from wiki.models.directory import Directory


@DB
class DirectoryDQ(DataSQLQuery):
    model = Directory

    def model_type_to_directory(self, model_type):
        model_type = int(model_type)
        dirctory_type = None
        if model_type in ITEM_H1:
            dirctory_type = DIRECTORY_TYPE.ITEM
        if model_type in DRUG_H1:
            dirctory_type = DIRECTORY_TYPE.DRUG
        if model_type in MATERIAL_H1:
            dirctory_type = DIRECTORY_TYPE.MATERIAL
        if model_type in MACHINE_H1:
            dirctory_type = DIRECTORY_TYPE.MACHINE
        return dirctory_type

    def update(self, updates, **kwargs):
        if updates.get('model_type'):
            updates['dirctory_type'] = self.model_type_to_directory(updates['model_type'])
        return super(DirectoryDQ.sqlquery, self).update(updates, **kwargs)

    def create(self, **kwargs):
        kwargs['dirctory_type'] = self.model_type_to_directory(kwargs['model_type'])
        return super(DirectoryDQ.sqlquery, self).create(**kwargs)
