import math
import logging
import datetime
from tracemalloc import start

from django.core.management import BaseCommand
from django.db import connection

from pingcap_sso.models import IPLocation
from pingcap_sso.lib.ip2location import bulk_lookup

logger = logging.getLogger(__name__)


class Command(BaseCommand):

    def add_arguments(self, parser):
        parser.add_argument(
            '--type',
            required=True,
            choices=['manual', 'daily'],
            help='type of how to sync ip location from track_events',
        )
        parser.add_argument(
            '--start_at',
            required=False,
            type=lambda s: datetime.datetime.strptime(s, '%Y-%m-%d'),
            help='sync ip from track_events.ip after this date, e.g. 2019-01-01'
        )
        parser.add_argument(
            '--end_at',
            required=False,
            type=lambda s: datetime.datetime.strptime(s, '%Y-%m-%d'),
            help='sync ip from track_events.ip before this date, e.g. 2019-01-01'
        )

    def handle(self, *args, **options):
        sync_type = options['type']
        start_at = options['start_at']
        end_at = options['end_at']

        if sync_type == 'daily':
            end_at = datetime.datetime.now()
            start_at = end_at.replace(hour=0, minute=0, microsecond=0) - datetime.timedelta(days=2)
            self.sync_ip_from_track_events(start_at, end_at)
        elif sync_type == 'manual':
            if not (start_at and end_at):
                logger.error('--start_at and --end_at are required when sync_type is manual')
                return
            if start_at > end_at:
                logger.error('--start_at must be less than --end_at')
                return

            days_step = 7
            while start_at < end_at:
                _next_end_at = start_at + datetime.timedelta(days=days_step)
                if _next_end_at > end_at:
                    _next_end_at = end_at
                self.sync_ip_from_track_events(start_at, _next_end_at)
                start_at = _next_end_at
        else:
            logger.error('invalid sync_type')
            return

        ip_qs = IPLocation.objects.filter(city__isnull=True)
        total_count = ip_qs.count()
        bulk_update_step = 500
        for i in range(math.ceil(total_count / bulk_update_step)):
            temp_ip_qs = ip_qs[0:bulk_update_step]
            ip_list = list(temp_ip_qs.values_list('ip', flat=True))
            ip_lookup_results = bulk_lookup(ip_list)
            for ip_instance in temp_ip_qs:
                ip_to_location = ip_lookup_results.get(ip_instance.ip)
                if ip_to_location:
                    ip_instance.country = ip_to_location['country_name'][:32]
                    ip_instance.region = ip_to_location['region'][:32]
                    ip_instance.city = ip_to_location['city'][:32]
                else:
                    ip_instance.city = 'unknown'
                ip_instance.save()

    def sync_ip_from_track_events(self, start_at, end_at):
        with connection.cursor() as cursor:
            cursor.execute("""
                INSERT IGNORE INTO pingcap_ip_location (ip)
                SELECT distinct ip
                FROM track_events
                WHERE
                    event_time > %s and event_time < %s and
                    ip is not null and ip <> ''
            """, [start_at, end_at])