import os, time, json, tarfile, tempfile, shutil
import requests

OUT_DIR = os.getenv("OUT_DIR", "/data")
LICENSE_KEY = os.getenv("MAXMIND_LICENSE_KEY", "").strip()
PDB_API_KEY = os.getenv("PDB_API_KEY", "").strip()
PDB_BASE = os.getenv("PDB_BASE", "https://www.peeringdb.com")
INFO_TYPE = os.getenv("PDB_INFO_TYPE", "Educational/Research")
TIMEOUT = int(os.getenv("HTTP_TIMEOUT", "30"))
LIMIT = int(os.getenv("PDB_LIMIT", "250"))

def atomic_replace(src_path: str, dst_path: str) -> None:
    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
    tmp = dst_path + ".tmp"
    shutil.copyfile(src_path, tmp)
    os.replace(tmp, dst_path)
    os.chmod(dst_path, 0o644)

def download_maxmind_mmdb() -> None:
    if not LICENSE_KEY:
        raise RuntimeError("MAXMIND_LICENSE_KEY missing")

    # Offizieller GeoLite2 Download-Mechanismus per license_key + edition_id
    url = (
        "https://download.maxmind.com/app/geoip_download"
        f"?edition_id=GeoLite2-ASN&license_key={LICENSE_KEY}&suffix=tar.gz"
    )

    with tempfile.TemporaryDirectory() as td:
        tgz = os.path.join(td, "GeoLite2-ASN.tar.gz")
        r = requests.get(url, timeout=TIMEOUT)
        if r.status_code == 429:
            existing = os.path.join(OUT_DIR, "GeoLite2-ASN.mmdb")
            if os.path.exists(existing):
                print("[warn] MaxMind rate limited (429); keeping existing mmdb")
                return
            raise RuntimeError("MaxMind rate limited (429) and no existing mmdb")
        r.raise_for_status()
        with open(tgz, "wb") as f:
            f.write(r.content)

        mmdb_found = None
        with tarfile.open(tgz, "r:gz") as tar:
            for member in tar.getmembers():
                if member.name.endswith("GeoLite2-ASN.mmdb"):
                    tar.extract(member, path=td)
                    mmdb_found = os.path.join(td, member.name)
                    break

        if not mmdb_found or not os.path.exists(mmdb_found):
            raise RuntimeError("GeoLite2-ASN.mmdb not found in archive")

        atomic_replace(mmdb_found, os.path.join(OUT_DIR, "GeoLite2-ASN.mmdb"))

def pdb_headers():
    if not PDB_API_KEY:
        return {"Accept": "application/json"}
    # PeeringDB API Key (optional)
    return {"Accept": "application/json", "Authorization": f"Api-Key {PDB_API_KEY}"}

def fetch_pdb_page(skip: int, info_type: str):
    url = f"{PDB_BASE}/api/net"
    params = {
        "info_type": info_type,
        "limit": LIMIT,
        "skip": skip,
        "fields": "asn,status,info_type",
    }
    r = requests.get(url, params=params, headers=pdb_headers(), timeout=TIMEOUT)
    r.raise_for_status()
    j = r.json()
    return j.get("data", [])

def update_nren_asns() -> str:
    info_types = [INFO_TYPE]
    # Alternate label seen in PeeringDB deployments.
    if INFO_TYPE != "Research and Education":
        info_types.append("Research and Education")
    if INFO_TYPE != "Educational/Research":
        info_types.append("Educational/Research")

    asns = set()
    used_info_type = INFO_TYPE
    for info_type in info_types:
        asns.clear()
        skip = 0
        while True:
            data = fetch_pdb_page(skip, info_type)
            for obj in data:
                if obj.get("status") != "ok":
                    continue
                asn = obj.get("asn")
                if isinstance(asn, int) and asn > 0:
                    asns.add(asn)
            if len(data) < LIMIT:
                break
            skip += LIMIT
            time.sleep(1.1)  # sehr konservativ
        if asns:
            used_info_type = info_type
            break

    if not asns:
        print(f"[warn] no ASNs found for info_type(s)={info_types}")

    out_txt = os.path.join(OUT_DIR, "nren_asns.txt")
    with tempfile.NamedTemporaryFile("w", delete=False, dir=OUT_DIR) as f:
        for a in sorted(asns):
            f.write(f"{a}\n")
        tmp_path = f.name
    os.replace(tmp_path, out_txt)
    os.chmod(out_txt, 0o644)
    return used_info_type

def write_meta(info_type: str):
    meta = {
        "updated_at_unix": int(time.time()),
        "info_type": info_type,
        "pdb_base": PDB_BASE,
    }
    with open(os.path.join(OUT_DIR, "metadata.json"), "w") as f:
        json.dump(meta, f, indent=2)
    os.chmod(os.path.join(OUT_DIR, "metadata.json"), 0o644)

def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    download_maxmind_mmdb()
    used_info_type = update_nren_asns()
    write_meta(used_info_type)
    print("[ok] updated mmdb + nren_asns")

if __name__ == "__main__":
    main()
