# Copyright (c) 2007, Justin Bronn
# All rights reserved.
#
# Released under New BSD License
#
"""
  These scripts are used to import the MaxMind(R) GeoIP Lite CSV files.

  In order to save memory during import ensure DEBUG=False in your settings.
"""
from models import Country, CountryBlock, Location, LocationBlock
from django.contrib.gis.geos import Point
from csv import reader
import sys

def country_import(csv_file):

    fh = open(csv_file)
    table = reader(fh)

    header = table.next()

    for startip, endip, ipfrom, ipto, country, country_name in table:
        cntry, created = Country.objects.get_or_create(name=country_name, code=country)
        if created: print 'Created: %s' % cntry
        block = CountryBlock(ipto=ipto, ipfrom=ipfrom, startip=startip, endip=endip, country=cntry)
        block.save()

    fh.close()
    del table

def location_import(loc_csv, block_csv):

    if loc_csv:
        # First, importing from the Location CSV file.
        fh = open(loc_csv)
        table = reader(fh)
        header = table.next()

        # Caching the countries table in memory
        countries = dict((m.code, m) for m in Country.objects.all())

        i = 0
        for locid, cntry, reg, cty, postal, lat, lon, dma, area in table:
            pnt = Point(float(lon), float(lat))

            # The points for the countries are in the first 244 entries --
            #  pulling these out and updating the points
            if int(locid) < 244:
                try:
                    country = Country.objects.get(code=cntry)
                    country.point = pnt
                except:
                    country = Country(code=cntry, point=pnt)
                country.save()
                countries[cntry] = country # updating the country dictionary
            else:
                country = countries[cntry]

            # region and city
            region = reg.decode('UTF-8', 'ignore')
            city = cty.decode('UTF-8', 'ignore')

            # Constructing the Location
            loc = Location(locid=locid, country=country, region=region, city=city,
                           postalcode=postal, point=pnt, dmacode=dma, areacode=area)
            loc.save()
            i += 1
            if i % 10000 == 0: print 'Saved %d Locations so far ...' % i
        fh.close()
        del table
        del countries

    if block_csv:
        # Second, importing from the Location IP block CSV file
        fh = open(block_csv)
        table = reader(fh)
        header = table.next()

        # This will take a little bit... and ~200+MB of RAM
        print 'Caching Location table...',
        sys.stdout.flush()
        locations = dict((m.locid, m) for m in Location.objects.all())
        print 'DONE.'

        i = 0
        for ipfrom, ipto, locid in table:
            loc = locations[int(locid)] # pulling location from our cached table (less expensive than Location.objects.get())
            loc_block = LocationBlock(location=loc, ipfrom=ipfrom, ipto=ipto)
            loc_block.save()
            i += 1
            if i % 10000 == 0: print 'Saved %d Location Blocks so far ...' % i

        fh.close()
        del table
        del locations