"""OFFLINE ASN lookup + provider branding (DB-IP ASN Lite)."""
import bisect
import csv
import ipaddress
import json
import os
import sqlite3

DEFAULT_MAP = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                           "asn_brand_map.json")


def load_brand_map(path=None):
    try:
        with open(path or DEFAULT_MAP, encoding="utf-8") as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        data = {"items": [], "defaults": {"brand": "未知", "type": "unknown",
                                          "icon": "❓"}}
    return data


def classify_org(org, brand_map=None):
    brand_map = brand_map or load_brand_map()
    text = (org or "").lower()
    for item in brand_map.get("items", []):
        if any(keyword in text for keyword in item.get("keywords", [])):
            return item["brand"], item["type"], item["icon"]
    fallback = brand_map.get("defaults") or {}
    return (fallback.get("brand", "未知"), fallback.get("type", "unknown"),
            fallback.get("icon", "❓"))


def build_asn_db(csv_path, out_path):
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    conn = sqlite3.connect(out_path)
    conn.execute("DROP TABLE IF EXISTS asn")
    conn.execute("CREATE TABLE asn (start INTEGER PRIMARY KEY, end INTEGER, "
                 "asn TEXT, org TEXT)")
    batch, total = [], 0

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

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


class AsnDB:
    def __init__(self, path, brand_map_path=None):
        self.conn = sqlite3.connect(path, check_same_thread=False)
        self.brand_map = load_brand_map(brand_map_path)
        self._prefix = None
        with self.conn:
            try:
                self.conn.execute("PRAGMA query_only=ON")
            except sqlite3.Error:
                pass

    def lookup(self, ip):
        try:
            value = int(ipaddress.IPv4Address(ip))
        except ValueError:
            return None
        row = self.conn.execute(
            "SELECT end, asn, org FROM asn WHERE start <= ? "
            "ORDER BY start DESC LIMIT 1", (value,)).fetchone()
        if not row or row[0] < value:
            return None
        brand, kind, icon = classify_org(row[2], self.brand_map)
        return {"asn": row[1], "org": row[2], "brand": brand,
                "type": kind, "icon": icon}

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


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