import csv
from django.core.management.base import BaseCommand
from prag.server.models import Proxy


class Command(BaseCommand):  # pragma: no cover
    help = 'Import list of proxies from a CSV file generated by proxy crawler.'

    def add_arguments(self, parser):
        parser.add_argument(
            '--model', '-m', dest='model', help='Model to import',
            default='proxy', choices=['proxy']
        )

        parser.add_argument(
            'filename', nargs=1, help='CSV file'
        )

    def handle(self, *args, **options):
        model, filename = options['model'], options['filename'][0]

        if model == 'proxy':
            self.import_proxy(filename)

    def import_proxy(self, filename):
        Proxy.objects.all().delete()
        with open(filename) as csvfile:
            reader = csv.DictReader(csvfile)
            count = 0
            for row in reader:
                Proxy.from_json(row)
                count += 1
        print('%d proxies imported' % count)
