from banapedia.wapi.WikipediaQuery import BlockQuery
from datetime import datetime
import pygeoip

__author__ = 'pacien'


GEOIP_FILE = "/usr/share/GeoIP/GeoIP.dat"
geoip = pygeoip.GeoIP(GEOIP_FILE)

ISO_TIMESTAMP = "%Y-%m-%dT%H:%M:%SZ"


class Ban:
    def __init__(self, ip, start, end):
        self.ip = ip
        self.start = start
        self.end = end
        self.country_code = None

    def get_duration(self):
        return (self.end - self.start).days

    def get_country_code(self):
        if self.country_code is not None:
            return self.country_code

        country_code = ""

        try:
            country_code = geoip.country_code_by_addr(self.ip).lower()
        except pygeoip.GeoIPError:
            print("[ERROR]", "Could not determine country for ip", self.ip)

        self.country_code = country_code
        return country_code


def map_ban(ban_dict):
    return Ban(
        ban_dict["user"],
        datetime.strptime(ban_dict["timestamp"], ISO_TIMESTAMP),
        datetime.strptime(ban_dict["expiry"], ISO_TIMESTAMP),
    )


def map_bans(ban_dict_list):
    ban_list = []
    for ban_dict in ban_dict_list:
        ban_list.append(map_ban(ban_dict))

    return ban_list


def fetch_multipart_ban_dict(n, query_limit):
    ban_dict_list = []
    n_fetched = 0
    continue_token = None

    print("[INFO]", "Fetching %d bans" % n)
    while n_fetched < n:
        to_fetch = min(query_limit, n - n_fetched)
        query = BlockQuery(
            bkprop=["user", "timestamp", "expiry"],
            bkshow=["temp", "ip"],
            limit=to_fetch,
            continue_token=continue_token,
        )
        results = query.fetch_result()
        ban_dict_list.extend(results["query"]["blocks"])
        continue_token = results["query-continue"]["blocks"]["bkcontinue"]
        n_fetched += to_fetch
        print("[INFO]", "Fetched %d over %d bans" % (n_fetched, n))

    print("[INFO]", "Bans fetching complete")
    return ban_dict_list