import csv
import lzma

from django.utils.timezone import now
from reprint import output

from jarbas.core.management.commands import LoadCommand
from jarbas.core.models import Reimbursement


class Command(LoadCommand):
    help = 'Load Serenata de Amor reimbursements dataset'

    def handle(self, *args, **options):
        self.started_at = now()
        self.path = options['dataset']

        existing = Reimbursement.objects.count()
        print('Starting with {:,} reimbursements'.format(existing))
        self.count = {'updated': 0, 'created': 0, 'skip': 0}

        if options.get('drop', False):
            self.drop_all(Reimbursement)

        with output() as status:
            status.change(self.status)
            self.create_or_update(self.reimbursements, status)

        self.mark_not_updated_reimbursements()

    @property
    def reimbursements(self):
        """Returns a Generator with a dict object for each row."""
        with lzma.open(self.path, mode='rt') as file_handler:
            for row in csv.DictReader(file_handler):
                yield self.serialize(row)

    def serialize(self, reimbursement):
        """Read the dict generated by DictReader and fix content types"""

        missing = ('probability', 'suspicions')
        for key in missing:
            reimbursement[key] = None

        rename = (
            ('subquota_number', 'subquota_id'),
            ('reimbursement_value_total', 'total_reimbursement_value')
        )
        for old, new in rename:
            reimbursement[new] = reimbursement[old]
            del reimbursement[old]

        integers = (
            'applicant_id',
            'batch_number',
            'congressperson_document',
            'congressperson_id',
            'document_id',
            'document_type',
            'installment',
            'month',
            'subquota_group_id',
            'subquota_id',
            'term',
            'term_id',
            'year'
        )
        for key in integers:
            reimbursement[key] = self.to_number(reimbursement[key], int)

        floats = (
            'document_value',
            'remark_value',
            'total_net_value',
            'total_reimbursement_value'
        )
        for key in floats:
            reimbursement[key] = self.to_number(reimbursement[key])

        reimbursement['issue_date'] = self.to_date(reimbursement['issue_date'])

        return reimbursement

    def create_or_update(self, reimbursements_as_dicts, status):
        for reimbursement in reimbursements_as_dicts:
            document_id = reimbursement.get('document_id')

            if not document_id:
                self.count['skip'] += 1
                status.change(self.status)
                continue

            _, created = Reimbursement.objects.update_or_create(
                document_id=document_id,
                defaults=reimbursement
            )

            key = 'created' if created else 'updated'
            self.count[key] += 1
            status.change(self.status)

    def mark_not_updated_reimbursements(self):
        qs = Reimbursement.objects.filter(last_update__lt=self.started_at)
        qs.update(available_in_latest_dataset=False)

    @property
    def status(self):
        label = '{}s'.format(self.get_model_name(Reimbursement)).lower()
        total = sum(map(self.count.get, self.count.keys()))
        output = [
            'Processed: {} lines'.format(total),
            'Updated: {} {}'.format(self.count['updated'], label),
            'Created: {} {}'.format(self.count['created'], label),
            'Skip: {} {}'.format(self.count['skip'], label),
        ]
        return output
