from ip2Region import Ip2Region


def ip_search(ip, region_type, algorithm='btree'):
    db_file = './ip2region.db'
    searcher = Ip2Region(dbfile=db_file)
    try:
        if algorithm == "binary":
            data = searcher.binarySearch(ip.strip())
        elif algorithm == "memory":
            data = searcher.memorySearch(ip.strip())
        else:
            data = searcher.btreeSearch(ip.strip())
        region = data['region'].decode('utf-8').split('|')
        if region_type == 'country':
            return region[0].strip()
        if region_type == 'province':
            return region[2].strip()
        if region_type == 'city':
            return region[3].strip()
    except Exception as e:
        # print(e)
        return None


def region2code(region):
    code_file = './ip_code.txt'
    code_map = {}
    with open(code_file, 'r', encoding='utf-8') as rf:
        lines = rf.readlines()
    for line in lines:
        seg_list = line.split('|')
        code_map[seg_list[0].strip()] = seg_list[1].strip()
    # print(code_map)
    code = code_map.get(region, None)
    # print(code)
    return code


if __name__ == '__main__':
    # ip = '101.105.35.57'
    # res = ip_search(ip=ip, region_type='city', algorithm='binary')
    # print(res)
    region2code('河北省')