#!/usr/bin/env python3
"""Reassemble split GeoIP database chunks into .mmdb.gz files, then decompress.

Chunks are stored in data/geoip_chunks/ to work around Git push size limits.
Run this once after cloning the repo, or let geoip.py auto-decompress at startup.
"""
import hashlib
import gzip
import shutil
from pathlib import Path

DATA_DIR = Path(__file__).parent.parent / 'data'
CHUNKS_DIR = DATA_DIR / 'geoip_chunks'

CHECKSUMS = {
    'GeoLite2-City.mmdb.gz': 'a0b0fd27d587ab5caf1436bdbe8ab0a4',
    'GeoLite2-ASN.mmdb.gz': '0b0fd27d587ab5caf1436bdbe8ab0a4',  # placeholder
}

def reassemble(name: str) -> Path | None:
    """Reassemble chunks for `name` into a single .gz file. Returns path or None."""
    prefix = name + '.gz.part'
    chunks = sorted(
        p for p in CHUNKS_DIR.iterdir()
        if p.name.startswith(prefix)
    )
    if not chunks:
        print(f'No chunks found for {name}.gz')
        return None

    out_path = DATA_DIR / (name + '.gz')
    print(f'Reassembling {len(chunks)} chunks → {out_path.name} …')
    with open(out_path, 'wb') as out:
        for chunk in chunks:
            with open(chunk, 'rb') as f:
                shutil.copyfileobj(f, out)
    print(f'Done: {out_path.stat().st_size:,} bytes')
    return out_path

def decompress(name: str) -> Path | None:
    """Decompress .mmdb.gz → .mmdb if not already present."""
    gz_path = DATA_DIR / (name + '.gz')
    mmdb_path = DATA_DIR / name
    if mmdb_path.exists():
        print(f'{name} already decompressed')
        return mmdb_path
    if not gz_path.exists():
        print(f'{name}.gz not found')
        return None
    print(f'Decompressing {gz_path.name} → {mmdb_path.name} ({gz_path.stat().st_size:,} bytes) …')
    with gzip.open(gz_path, 'rb') as f_in:
        with open(mmdb_path, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out, length=16 * 1024 * 1024)
    print(f'Done: {mmdb_path.stat().st_size:,} bytes')
    return mmdb_path

def main():
    for name in ['GeoLite2-City.mmdb', 'GeoLite2-ASN.mmdb']:
        gz_path = DATA_DIR / (name + '.gz')
        if not gz_path.exists():
            reassemble(name)
        decompress(name)
    print('GeoIP databases ready.')

if __name__ == '__main__':
    main()
