"""Offline GeoIP lookup backed by the free DB-IP Lite CSV dumps.

The city dump has no country name, so a small code -> name table is shipped
here; anything unknown falls back to the ISO code.
"""
import csv
import ipaddress
import os
import sqlite3

COUNTRY_CN = {
    "CN": "中国", "HK": "中国香港", "TW": "中国台湾", "MO": "中国澳门",
    "US": "美国", "CA": "加拿大", "MX": "墨西哥", "BR": "巴西", "AR": "阿根廷",
    "CL": "智利", "CO": "哥伦比亚", "PE": "秘鲁", "VE": "委内瑞拉",
    "GB": "英国", "IE": "爱尔兰", "FR": "法国", "DE": "德国", "NL": "荷兰",
    "BE": "比利时", "LU": "卢森堡", "CH": "瑞士", "AT": "奥地利", "IT": "意大利",
    "ES": "西班牙", "PT": "葡萄牙", "GR": "希腊", "PL": "波兰", "CZ": "捷克",
    "SK": "斯洛伐克", "HU": "匈牙利", "RO": "罗马尼亚", "BG": "保加利亚",
    "HR": "克罗地亚", "RS": "塞尔维亚", "SI": "斯洛文尼亚", "LT": "立陶宛",
    "LV": "拉脱维亚", "EE": "爱沙尼亚", "FI": "芬兰", "SE": "瑞典",
    "NO": "挪威", "DK": "丹麦", "IS": "冰岛", "UA": "乌克兰", "BY": "白俄罗斯",
    "RU": "俄罗斯", "TR": "土耳其", "IL": "以色列", "SA": "沙特阿拉伯",
    "AE": "阿联酋", "QA": "卡塔尔", "KW": "科威特", "IR": "伊朗", "IQ": "伊拉克",
    "EG": "埃及", "ZA": "南非", "NG": "尼日利亚", "KE": "肯尼亚", "MA": "摩洛哥",
    "ET": "埃塞俄比亚", "IN": "印度", "PK": "巴基斯坦", "BD": "孟加拉国",
    "LK": "斯里兰卡", "NP": "尼泊尔", "MM": "缅甸", "KH": "柬埔寨", "LA": "老挝",
    "TH": "泰国", "VN": "越南", "MY": "马来西亚", "SG": "新加坡", "ID": "印度尼西亚",
    "PH": "菲律宾", "JP": "日本", "KR": "韩国", "MN": "蒙古", "KZ": "哈萨克斯坦",
    "UZ": "乌兹别克斯坦", "GE": "格鲁吉亚", "AM": "亚美尼亚", "AZ": "阿塞拜疆",
    "AU": "澳大利亚", "NZ": "新西兰", "ZZ": "未知",
}

CN_REGION = {
    "Beijing": "北京", "Tianjin": "天津", "Hebei": "河北", "Shanxi": "山西",
    "Inner Mongolia": "内蒙古", "Liaoning": "辽宁", "Jilin": "吉林",
    "Heilongjiang": "黑龙江", "Shanghai": "上海", "Jiangsu": "江苏",
    "Zhejiang": "浙江", "Anhui": "安徽", "Fujian": "福建", "Jiangxi": "江西",
    "Shandong": "山东", "Henan": "河南", "Hubei": "湖北", "Hunan": "湖南",
    "Guangdong": "广东", "Guangxi": "广西", "Hainan": "海南", "Chongqing": "重庆",
    "Sichuan": "四川", "Guizhou": "贵州", "Yunnan": "云南", "Tibet": "西藏",
    "Shaanxi": "陕西", "Gansu": "甘肃", "Qinghai": "青海", "Ningxia": "宁夏",
    "Xinjiang": "新疆", "Taiwan": "台湾", "Hong Kong": "香港", "Macau": "澳门",
}


def flag(cc):
    if not cc or len(cc) != 2 or not cc.isalpha() or cc == "ZZ":
        return ""
    return "".join(chr(0x1F1E6 + ord(char) - 65) for char in cc.upper())


def country_name(cc):
    return COUNTRY_CN.get(cc, cc or "")


def region_name(cc, region):
    if cc == "CN" and region:
        return CN_REGION.get(region, region)
    return region or ""


def build_geo_db(city_csv, out_db, country_csv=None):
    """Import the DB-IP city CSV (IPv4 only) into a small SQLite database."""
    os.makedirs(os.path.dirname(out_db) or ".", exist_ok=True)
    conn = sqlite3.connect(out_db)
    conn.execute("DROP TABLE IF EXISTS geo")
    conn.execute("CREATE TABLE geo (start INTEGER PRIMARY KEY, end INTEGER, "
                 "cc TEXT, region TEXT, city TEXT)")
    batch = []
    total = 0

    def flush():
        nonlocal batch
        if batch:
            conn.executemany("INSERT OR REPLACE INTO geo VALUES (?,?,?,?,?)", batch)
            batch = []

    with open(city_csv, newline="", encoding="utf-8", errors="replace") as fh:
        for row in csv.reader(fh):
            if len(row) < 6:
                continue
            start, end, cc = row[0], row[1], row[3]
            if ":" in start or ":" in end or not cc:
                continue
            try:
                first = int(ipaddress.IPv4Address(start))
                last = int(ipaddress.IPv4Address(end))
            except ValueError:
                continue
            batch.append((first, last, cc, row[4], row[5]))
            total += 1
            if len(batch) >= 100000:
                flush()
                conn.commit()
    flush()
    conn.commit()
    conn.execute("PRAGMA journal_mode=WAL")
    conn.close()
    return total


class GeoIP:
    """Read-only lookup over the generated geo.db."""

    def __init__(self, path):
        if not path or not os.path.exists(path):
            raise FileNotFoundError(path)
        self.conn = sqlite3.connect("file:%s?mode=ro" % path, uri=True,
                                    check_same_thread=False)

    def lookup(self, ip):
        try:
            value = int(ipaddress.ip_address(ip))
        except ValueError:
            return None
        row = self.conn.execute(
            "SELECT end, cc, region, city FROM geo WHERE start <= ? "
            "ORDER BY start DESC LIMIT 1", (value,)).fetchone()
        if not row or row[0] < value:
            return None
        return {"country_code": row[1], "country_name": country_name(row[1]),
                "region": region_name(row[1], row[2]), "city": row[3] or ""}

    def close(self):
        try:
            self.conn.close()
        except sqlite3.Error:
            pass


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--city-csv", default="data/geo/dbip-city-lite-2026-10.csv")
    parser.add_argument("--out", default="data/geo/geo.db")
    args = parser.parse_args()
    count = build_geo_db(args.city_csv, args.out)
    print("imported %d IPv4 ranges -> %s" % (count, args.out))