bulk-target-7-april

import csv
import io
from collections import defaultdict, deque
from csv import DictReader

import graphene
from django.db import transaction
from graphene.types import InputObjectType
from django.core.exceptions import ValidationError

from ...core.mutations import ModelMutation, BaseMutation
from ...core.types import Upload
from ...core.types.common import TargetError
from ....account.models import User
from ....core.permissions import TargetPermissions
from ....partner.models import Partner
from ....target import models
from ....target.error_codes import TargetErrorCode
from dateutil import parser, relativedelta
from datetime import datetime as dt

from ....target.models import Achievement, TargetUser, PartnerTargetSales
from ...utils import check_if_attempted_from_valid_parent
from ...target.types import PartnerTargetSales as PartnerTargetSalesType


class CreateTargetInput(InputObjectType):
    user = graphene.ID(description="ID of the Agent.", name="user", required=True)
    month = graphene.Date(required=True, description="Target month(YYYY-MM)")
    target_orders = graphene.Int(
        description="Number of orders for the target.", required=True
    )
    target_sales = graphene.Int(
        description="Amount of sales for the target.", required=True
    )


class CreateTarget(ModelMutation):
    class Arguments:
        input = CreateTargetInput(
            required=True, description="Fields required to create a target."
        )

    class Meta:
        description = "Creates a new target of an agent."
        model = models.Target
        permissions = (TargetPermissions.MANAGE_TARGETS,)
        error_type_class = TargetError
        error_type_field = "target_errors"

    @classmethod
    def clean_input(cls, info, instance, data):
        user = cls.get_node_or_error(info, data["user"])
        requestor = info.context.user
        if requestor.groups.first().name in ["dco", "dcm", "cm"]:
            check_if_attempted_from_valid_parent(requestor, user)

        cleaned_input = super().clean_input(info, instance, data)
        if cleaned_input['target_orders'] < 1:
            raise ValidationError(
                {
                    "target_orders": ValidationError(
                        "Target order should be positive number",
                        code=TargetErrorCode.INVALID,
                    )
                }
            )

        if cleaned_input['target_sales'] < 1:
            raise ValidationError(
                {
                    "target_sales": ValidationError(
                        "Target sales should be positive number",
                        code=TargetErrorCode.INVALID,
                    )
                }
            )

        user = cleaned_input['user']
        month = parser.parse(str(cleaned_input['month'])).replace(day=1)
        cleaned_input['month'] = month
        target = models.Target.objects.filter(user=user, month=month)

        if target:
            raise ValidationError(
                {
                    "month": ValidationError(
                        "Target for specified month already exists.",
                        code=TargetErrorCode.ALREADY_EXISTS,
                    )
                }
            )

        return cleaned_input

    @classmethod
    def save(cls, info, instance, cleaned_input):
        instance.save()
        user = cleaned_input['user']
        month = cleaned_input['month']
        achievement = Achievement.objects.filter(user=user, month=month).first()
        if achievement:
            achievement.target = instance
            achievement.save()


class UpdateTargetInput(graphene.InputObjectType):
    target_orders = graphene.Int(
        description="Number of orders for the target.", required=False
    )
    target_sales = graphene.Int(
        description="Amount of sales for the target.", required=False
    )


class UpdateTarget(ModelMutation):
    class Arguments:
        id = graphene.ID(required=True, description="ID of a target to update.")
        input = UpdateTargetInput(
            required=True, description="Fields required to update a target."
        )

    class Meta:
        description = "Updates a new target of an agent."
        model = models.Target
        permissions = (TargetPermissions.MANAGE_TARGETS,)
        error_type_class = TargetError
        error_type_field = "target_errors"

    @classmethod
    def clean_input(cls, info, instance, data):
        target_month = instance.month.strftime("%Y-%m")
        current_month = dt.now().date().strftime("%Y-%m")
        if target_month < current_month:
            raise ValidationError(
                {
                    "id": ValidationError(
                        "This item can not be modified.",
                        code=TargetErrorCode.INVALID
                    )
                }
            )
        target_orders = data.get("target_orders")
        target_sales = data.get("target_sales")

        if not target_sales and not target_orders:
            raise ValidationError(
                {
                    "target_sales": ValidationError(
                        "Target sales can not be empty",
                        code=TargetErrorCode.INVALID
                    ),
                    "target_orders": ValidationError(
                        "Target orders can not be empty",
                        code=TargetErrorCode.INVALID
                    )
                }
            )

        if target_orders < 1:
            raise ValidationError(
                {
                    "target_orders": ValidationError(
                        "Target order should be positive number",
                        code=TargetErrorCode.INVALID
                    )
                }
            )

        if target_sales < 1:
            raise ValidationError(
                {
                    "target_sales": ValidationError(
                        "Target sales should be positive number",
                        code=TargetErrorCode.INVALID,
                    )
                }
            )

        return super().clean_input(info, instance, data)


class CreateTargetsInBulk(BaseMutation):
    partner_target_sales = graphene.List(PartnerTargetSalesType, description="List of partner target sales.")
    graph = defaultdict(list)

    class Arguments:
        file = Upload(
            required=True,
            description="A csv file that contains agent email, partner id, target amount to set targets in bulk.",
        )

    class Meta:
        description = "Upload file(csv) to create targets in bulk."
        error_type_class = TargetError
        error_type_field = "target_errors"
        permissions = (TargetPermissions.MANAGE_TARGETS,)

    @classmethod
    def get_parents_list(self, source, include_self=False):
        if len(self.graph) == 0:
            all_users = User.objects.all()
            for user in all_users:
                if user.parent_id:
                    self.graph[user.id].append(user.parent_id)

        parents_list = []
        visited = set()
        queue = deque()

        visited.add(source)
        queue.append(source)

        while queue:
            s = queue.popleft()
            for neighbour in self.graph[s]:
                if neighbour not in visited:
                    parents_list.append(neighbour)
                    visited.add(neighbour)
                    queue.append(neighbour)

        if include_self:
            parents_list.append(source)

        return parents_list

    @classmethod
    def perform_mutation(cls, _root, info, file):
        with transaction.atomic():
            csv_file = info.context.FILES.get(file)

            if not csv_file.name.endswith(".csv"):
                raise ValidationError(
                    {
                        "target_sales": ValidationError(
                            "The file provided is not in csv format.",
                            code=TargetErrorCode.INVALID,
                        )
                    }
                )

            data_set = csv_file.read().decode("UTF-8")
            io_string = io.StringIO(data_set)

            csv_reader = csv.reader(io_string, delimiter=',')
            header = next(csv_reader)
            no_of_columns = len(header)

            agents = {}
            partners = {}
            target_user_objects_cache = {}
            partner_target_sales_cache = {}

            all_agents = User.objects.filter(groups__name='agent')
            all_partners = Partner.objects.all()

            for agent in all_agents:
                agents[agent.email] = agent

            for partner in all_partners:
                partners[partner.partner_id] = partner

            target_user_list = []
            partner_target_sales_list_new = []
            partner_target_sales_list_existing = []

            partner_target_sales_hash_new = {}
            partner_target_sales_hash_existing = {}
            raw_data = []
            current_month = dt.now().replace(day=1).date()
            agent_ids = set()
            agent_months = defaultdict(set)

            for row in csv_reader:
                raw_data.append(
                    {
                        'row': row,
                        'target_user_objects_list_new': [],
                        'target_user_objects_list_existing': [],

                    }
                )

            target_users_all = TargetUser.objects.filter(created__gte=current_month)

            for item in target_users_all:
                hash_val = str(item.user_id) + "-" + str(item.month)
                target_user_objects_cache[hash_val] = item

            for item in raw_data:
                email = item['row'][0]
                month = item['row'][1]
                date_object = parser.parse(str(month)).replace(day=1).date()
                month = str(date_object)
                agent = agents.get(email, None)

                if date_object < current_month:
                    raise ValidationError(
                        {
                            "month": ValidationError(
                                "%s is invalid, current or future month is allowed only." % item['row'][1],
                                code=TargetErrorCode.INVALID,
                            )
                        }
                    )

                if agent is None:
                    raise ValidationError(
                        {
                            "agent": ValidationError(
                                "%s email is invalid." % email,
                                code=TargetErrorCode.INVALID,
                            )
                        }
                    )

                users_list = cls.get_parents_list(agent.id, True)
                agent_ids.add(agent.id)
                agent_months[agent.id].add(month)

                for user_id in users_list:
                    hash_val = str(user_id) + "-" + str(month)
                    target_user_item = target_user_objects_cache.get(hash_val, None)
                    if target_user_item is None:
                        target_user_item = TargetUser(
                            user_id=user_id,
                            month=month
                        )
                        target_user_list.append(target_user_item)
                        item['target_user_objects_list_new'].append(target_user_item)

                        target_user_objects_cache[hash_val] = target_user_item  # added in cache
                    else:
                        if target_user_item.id is None:
                            item['target_user_objects_list_new'].append(target_user_item)
                        else:
                            item['target_user_objects_list_existing'].append(target_user_item)

            TargetUser.objects.bulk_create(target_user_list)

            partner_target_sales_all = PartnerTargetSales.objects\
                .filter(created__gte=current_month).prefetch_related('target_user')

            for item in partner_target_sales_all:
                _user_id = item.target_user.user_id
                _month = str(item.target_user.month)

                if _user_id in agent_ids and _month in agent_months[_user_id]:
                    raise ValidationError(
                        {
                            "agent": ValidationError(
                                "%s already has a target for %s." % (item.target_user.user.email, _month),
                                code=TargetErrorCode.INVALID,
                            )
                        }
                    )
                hash_val = str(item.target_user_id) + "-" + str(item.partner_id)
                partner_target_sales_cache[hash_val] = item

            for item in raw_data:
                target_user_objects_list_new = item['target_user_objects_list_new']
                target_user_objects_list_existing = item['target_user_objects_list_existing']
                target_user_objects_list_new.extend(target_user_objects_list_existing)

                if len(target_user_objects_list_new) > 0:
                    for target_user_obj in target_user_objects_list_new:
                        index = 2
                        while index + 1 < no_of_columns:
                            if len(item['row'][index]) == 0:
                                index += 2
                                continue

                            partner_id = item['row'][index]
                            target_sales = int(item['row'][index + 1])

                            partner_obj = partners.get(partner_id, None)

                            if partner_obj is None:
                                raise ValidationError(
                                    {
                                        "partner": ValidationError(
                                            "%s partner id is invalid." % partner_id,
                                            code=TargetErrorCode.INVALID,
                                        )
                                    }
                                )

                            hash_val = str(target_user_obj.id) + "-" + str(partner_obj.id)
                            partner_target_sales_item = partner_target_sales_cache.get(hash_val, None)

                            if partner_target_sales_item is None:

                                partner_target_sales_item = PartnerTargetSales(
                                    target_user=target_user_obj,
                                    partner=partner_obj,
                                    target_sales=target_sales
                                )
                                partner_target_sales_hash_new[hash_val] = partner_target_sales_item
                                partner_target_sales_cache[hash_val] = partner_target_sales_item

                            else:
                                partner_target_sales_item.target_sales += target_sales

                                if partner_target_sales_item.id is None:
                                    partner_target_sales_hash_new[hash_val] = partner_target_sales_item
                                else:
                                    partner_target_sales_hash_existing[hash_val] = partner_target_sales_item

                            index += 2

            for key, val in partner_target_sales_hash_new.items():
                partner_target_sales_list_new.append(val)

            for key, val in partner_target_sales_hash_existing.items():
                partner_target_sales_list_existing.append(val)

            PartnerTargetSales.objects.bulk_create(partner_target_sales_list_new)
            PartnerTargetSales.objects.bulk_update(partner_target_sales_list_existing, ['target_sales', 'updated'])

            partner_target_sales = PartnerTargetSales.objects.all()

            return CreateTargetsInBulk(partner_target_sales=partner_target_sales)