import os
from pathlib import Path

from celery import shared_task
from celery.utils.log import get_task_logger
from django.conf import settings
from django.contrib.auth.models import User
from django.core.management import call_command

from storage.export import export
from storage.models import DataStorage, get_storage_path

_log = get_task_logger(__name__)

MEDIA_ROOT = getattr(settings, 'MEDIA_ROOT')


@shared_task
def export_task(export_type, user_pk, model_name, query_parameters):
    return export(export_type, user_pk, model_name, query_parameters)


@shared_task
def dump_db(pk, dump_models):
    # Add store record
    user = User.objects.get(pk=pk)
    _log.info(f'User [{user}] request db dump for models {dump_models}')
    store = DataStorage(type=DataStorage.Type.DB_DUMP,
                        format=DataStorage.Format.JSON,
                        user=user,
                        size=0,
                        file='name')
    store.save()
    file_name = get_storage_path(store, 'dump.json')

    fpath = os.path.join(MEDIA_ROOT, file_name)

    dirname = os.path.dirname(fpath)
    Path(dirname).mkdir(parents=True, exist_ok=True)

    with open(fpath, 'w') as f:
        call_command('dumpdata', *dump_models, stdout=f)

    store.size = os.path.getsize(fpath)
    store.file = file_name
    store.save()

    _log.info(f'Dump ready in [{store.file.name}] with size of [{store.size}]')

    store.update_crc()

    return store.pk


def remove_vulnerabilities(data_dict):
    """
    Remove column vulnerabilities from data
    return: DataDict without vulnerabilities column
    """
    for item in data_dict:
        try:
            item.pop('vulnerabilities')
        except KeyError:
            continue

    return data_dict
