from __future__ import absolute_import

import socket
import idiokit
from idiokit import dns

from . import utils, transformation


def _parse_ip(string, families=(socket.AF_INET, socket.AF_INET6)):
    for family in families:
        try:
            return socket.inet_ntop(family, socket.inet_pton(family, string))
        except (ValueError, socket.error):
            pass
    return None


def _nibbles(ipv6, _hex="0123456789abcdef"):
    result = []
    for ch in socket.inet_pton(socket.AF_INET6, ipv6):
        num = ord(ch)
        result.append(_hex[num >> 4])
        result.append(_hex[num & 0xf])
    return result


def _split(txt_results, keys):
    results = set()

    for strings in txt_results:
        pieces = "".join(strings).split("|")
        decoded = map(lambda x: x.strip().decode("utf-8", "replace"), pieces)

        item_list = list()
        for key, value in zip(keys, decoded):
            if key is None:
                continue
            if value in ("", "-"):
                continue
            item_list.append((key, value))

        if not item_list:
            continue

        results.add(frozenset(item_list))

    return tuple(tuple(x) for x in results)


class ASNameLookup(object):
    _keys = (None, None, None, "as allocated", "as name")

    def __init__(self, resolver=None, cache_time=4 * 60 * 60, catch_error=True):
        self._resolver = resolver
        self._cache = utils.TimedCache(cache_time)
        self._catch_error = catch_error

    @idiokit.stream
    def lookup(self, asn):
        try:
            asn = int(asn)
        except ValueError:
            idiokit.stop(())

        results = self._cache.get(asn, None)
        if results is not None:
            idiokit.stop(results)

        try:
            txt_results = yield dns.txt(
                "AS{0}.asn.cymru.com".format(asn),
                resolver=self._resolver)
        except dns.DNSError:
            if self._catch_error:
                idiokit.stop(())
            else:
                raise

        results = _split(txt_results, self._keys)
        self._cache.set(asn, results)
        idiokit.stop(results)


class OriginLookup(object):
    _keys = ("asn", "bgp prefix", "cc", "registry", "bgp prefix allocated")

    def __init__(self, resolver=None, cache_time=4 * 60 * 60, catch_error=True):
        self._resolver = resolver
        self._cache = utils.TimedCache(cache_time)
        self._catch_error = catch_error

    @idiokit.stream
    def _lookup(self, cache_key, query):
        results = self._cache.get(cache_key, None)
        if results is not None:
            idiokit.stop(results)

        try:
            txt_results = yield dns.txt(query, resolver=self._resolver)
        except dns.DNSError:
            if self._catch_error:
                idiokit.stop(())
            else:
                raise

        results = []
        for result in _split(txt_results, self._keys):
            result_dict = dict(result)
            for asn in result_dict.get("asn", "").split():
                if not asn:
                    continue
                result_dict["asn"] = asn
                results.append(tuple(result_dict.iteritems()))

        self._cache.set(cache_key, tuple(results))
        idiokit.stop(results)

    @idiokit.stream
    def lookup(self, ip):
        ipv4 = _parse_ip(ip, families=[socket.AF_INET])
        if ipv4 is not None:
            prefix = ".".join(reversed(ipv4.split(".")))
            results = yield self._lookup(ipv4, prefix + ".origin.asn.cymru.com")
            idiokit.stop(results)

        ipv6 = _parse_ip(ip, families=[socket.AF_INET6])
        if ipv6 is not None:
            prefix = ".".join(reversed(_nibbles(ipv6)))
            results = yield self._lookup(ipv6, prefix + ".origin6.asn.cymru.com")
            idiokit.stop(results)

        idiokit.stop(())


class CymruWhois(object):
    def __init__(self, resolver=None, cache_time=4 * 60 * 60):
        self._origin_lookup = OriginLookup(resolver, cache_time)
        self._asname_lookup = ASNameLookup(resolver, cache_time)

    def _ip_values(self, event, keys):
        for key in keys:
            for value in event.values(key, parser=_parse_ip):
                yield value

    @idiokit.stream
    def augment(self, *ip_keys):
        while True:
            event = yield idiokit.next()
            if not ip_keys:
                values = event.values(parser=_parse_ip)
            else:
                values = self._ip_values(event, ip_keys)

            for ip in values:
                items = yield self.lookup(ip)
                for key, value in items:
                    event.add(key, value)

            yield idiokit.send(event)

    @idiokit.stream
    def lookup(self, ip):
        results = yield self._origin_lookup.lookup(ip)
        for result in results:
            result = dict(result)

            asn = result.get("asn", None)
            if asn is not None:
                infos = yield self._asname_lookup.lookup(asn)
                for info in infos:
                    result.update(info)
                    break
                idiokit.stop(tuple(result.items()))
        idiokit.stop(())


global_whois = CymruWhois()

augment = global_whois.augment
lookup = global_whois.lookup


class Handler(transformation.Handler):
    def __init__(self, ip_keys=[], *args, **keys):
        transformation.Handler.__init__(self, *args, **keys)

        self.ip_keys = tuple(ip_keys)

    def transform(self):
        return augment(*self.ip_keys)