""" Implements a IP -> Country lookup.
"""

import csv
import gzip
import re
#import ipaddress

from bisect import bisect_left
from collections import defaultdict

ip_parser = re.compile('^(\d+)\.(\d+)\.(\d+)\.(\d+)$')
def atonl(s):
    """ Convert a string IP to network byte order int.

    This is roughly 2x as fast as int(ipaddress.ip_address(s.decode('utf-8')))
    """
    result = re.match(ip_parser, s)
    if not result:
        return -1
    o1 = int(result.group(1))
    o2 = int(result.group(2))
    o3 = int(result.group(3))
    o4 = int(result.group(4))
    if o1 > 255 or o2 > 255 or o3 > 255 or o4 > 255:
        return -1
    return o1 << 24 | o2 << 16 | o3 << 8 | o4

class IPGeo(object):
    BAD_IP = 'BAD_IP'
    UNKNOWN_IP_RANGE = 'UNKNOWN_IP_RANGE'

    def __init__(self, path):
        """ Expects path to the file generated by create_ip_country_csv_gz().
        """
        print("Loading IP Country map...")
        with gzip.GzipFile(path) as f:
            reader = csv.reader(f)
            header = next(reader)
            assert header == ['first', 'last', 'country']
            self._ranges = [(int(row[0]), int(row[1]), row[2]) for row in reader]
        self._lasts = [row[1] for row in self._ranges]
        print("Loaded %d ranges." % len(self._ranges))

    def lookup_country(self, ip_str):
        if not ip_str:
            return self.BAD_IP
        ip = atonl(ip_str)
        if ip == -1:
            return self.BAD_IP
        index = bisect_left(self._lasts, ip)
        if index < len(self._ranges):
            entry = self._ranges[index]
            if ip >= entry[0] and ip <= entry[1]:
                return entry[2]
        return self.UNKNOWN_IP_RANGE


def create_ip_country_csv_gz(country_locations_path, country_blocks_path, out_path):
    geoid_to_country = {}
    with open(country_locations_path) as f:
        reader = csv.reader(f)
        header = next(reader)
        assert header[0] == 'geoname_id'
        assert header[4] == 'country_iso_code'
        for row in reader:
            geoid_to_country[row[0]] = row[4]

    ranges = []
    with open(country_blocks_path) as f:
        reader = csv.reader(f)
        header = next(reader)
        assert header[0] == 'network'
        # TODO(chris): Should I be using header[3] == represented_country_geoname_id instead?
        assert header[2] == 'registered_country_geoname_id'
        for row in reader:
            if not row[2]:
                continue
            network = ipaddress.ip_network(row[0].decode('utf-8'))
            ranges.append([int(network.network_address), int(network.broadcast_address), geoid_to_country[row[2]]])

        ranges = sorted(ranges, key=lambda x: x[0])

        with gzip.GzipFile(out_path, 'w') as f:
            writer = csv.writer(f)
            writer.writerow(['first', 'last', 'country'])
            for row in ranges:
                writer.writerow(row)
