#!/usr/bin/env python3
"""Stream a large OSM PBF and write a cropped copy for one or more bboxes.

Unlike `osmium extract -s complete_ways` this never holds more than one
region's worth of objects in memory (a few sequential streaming passes), so
it works on a full-country PBF on a RAM-constrained machine.

File format (osmium / OSM-binary):
    [uint32 BE: BlobHeader length][BlobHeader][Blob: datasize bytes]
    BlobHeader { string type=1; bytes indexdata=2; int32 datasize=3 }
    Blob { bytes raw=1; int32 raw_size=2; bytes zlib_data=3; bytes zstd_data=7 }
    the data field holds a PrimitiveBlock message.

Coordinate encoding: deg = (value + offset) * granularity / 1e9
(default granularity 100, offsets 0 => value / 1e7).

Usage: crop_pbf.py INPUT.pbf OUTPUT.pbf MINLON MINLAT MAXLON MAXLAT [BBOX2 ...]

Ways touching a box are kept whole; their nodes are kept even slightly
outside the box. Relations and node metadata are dropped (the foot profile
ignores turn restrictions).
"""
import sys, zlib, struct

# ---------------- varint / zigzag ----------------
def read_varint(b, i):
    r = 0; s = 0
    while True:
        x = b[i]; i += 1
        r |= (x & 0x7f) << s
        if not x & 0x80: return r, i
        s += 7

def read_sint(b, i):
    v, i = read_varint(b, i)
    return (~v >> 1) if v & 1 else (v >> 1), i

def packed_sint(b, i, end):
    out = []
    while i < end:
        v, i = read_varint(b, i)
        out.append((~v >> 1) if v & 1 else (v >> 1))
    return out, i

def packed_uint(b, i, end):
    out = []
    while i < end:
        v, i = read_varint(b, i)
        out.append(v)
    return out, i

def skip_field(b, i, end, wt, num=None):
    if wt == 0:
        _, i = read_varint(b, i)
    elif wt == 1:
        i += 4
    elif wt == 2:
        l, i = read_varint(b, i); i += l
    elif wt == 3:
        i += 8
    elif wt == 4:                       # deprecated start-group
        while i < end:
            f, i = read_varint(b, i)
            if (f & 7) == 6 and (f >> 3) == num:
                break
            i = skip_field(b, i, end, f & 7, f >> 3)
    elif wt == 5:
        i += 8
    elif wt == 6:
        pass                            # end-group
    else:
        raise ValueError('wire type %d' % wt)
    return i

# ---------------- writers ----------------
def varint(n):
    out = bytearray()
    while True:
        x = n & 0x7f; n >>= 7
        out.append(x | (0x80 if n else 0))
        if not n: return bytes(out)

def zig(n):
    return (n << 1) ^ (n >> 63)

def tag(num, wt=0):
    return varint((num << 3) | wt)

def f_varint(num, v):
    return tag(num) + varint(v)

def f_sint(num, v):
    return tag(num) + varint(zig(v))

def f_msg(num, payload):
    return tag(num, 2) + varint(len(payload)) + payload

def f_packed_sint(num, vals):
    body = b''.join(varint(zig(v)) for v in vals)
    return f_msg(num, body)

def f_packed_uint(num, vals):
    body = b''.join(varint(v) for v in vals)
    return f_msg(num, body)

# ---------------- PrimitiveBlock parsing ----------------
def parse_block(payload):
    """Return list of (group, strings, gran, lato, lono).
    group = {'nodes': [(id, lat, lon)], 'dense': [(id, lat, lon)], 'ways': [(id, refs, keys, vals)]}
    node coordinates are returned in raw PBF units (not yet decoded to degrees)."""
    i = 0; n = len(payload)
    strings = ['']
    gran = 100; lato = 0; lono = 0
    out = []
    while i < n:
        f, i = read_varint(payload, i); num = f >> 3; wt = f & 7
        if num == 1 and wt == 2:                                   # stringtable
            l, i = read_varint(payload, i)
            j = i; end = i + l; s = []
            while j < end:
                fj, j = read_varint(payload, j); fn = fj >> 3; fw = fj & 7
                if fn == 1 and fw == 2:
                    sl, j = read_varint(payload, j)
                    s.append(payload[j:j+sl].decode('utf-8', 'replace')); j += sl
                else:
                    j = skip_field(payload, j, end, fw)
            strings = s          # the table already starts with the reserved '' entry
            i = end
        elif num in (17, 19, 20) and wt == 0:
            v, i = read_varint(payload, i)
            if num == 17: gran = v
            elif num == 19: lato = v
            else: lono = v
        elif num == 2 and wt == 2:                                 # primitivegroup
            l, i = read_varint(payload, i)
            j = i; end = i + l
            group = {'nodes': [], 'dense': [], 'ways': []}
            while j < end:
                gj, j = read_varint(payload, j); gn = gj >> 3; gw = gj & 7
                if gn == 1 and gw == 2:                            # Node
                    nl, j = read_varint(payload, j); ne = j + nl
                    nid = 0; la = 0; lo = 0
                    while j < ne:
                        k, j = read_varint(payload, j); kn = k >> 3; kw = k & 7
                        if kn == 1 and kw == 0:
                            nid, j = read_sint(payload, j)
                        elif kn in (8, 9) and kw == 0:
                            v, j = read_sint(payload, j)
                            if kn == 8: la = v
                            else: lo = v
                        else:
                            j = skip_field(payload, j, ne, kw)
                    group['nodes'].append((nid, la, lo))
                elif gn == 2 and gw == 2:                          # DenseNodes
                    dl, j = read_varint(payload, j); de = j + dl
                    ids = lats = lons = None
                    while j < de:
                        k, j = read_varint(payload, j); kn = k >> 3; kw = k & 7
                        if kn in (1, 8, 9) and kw == 2:
                            kl, j = read_varint(payload, j)
                            arr, j = packed_sint(payload, j, j + kl)
                            if kn == 1: ids = arr
                            elif kn == 8: lats = arr
                            else: lons = arr
                        else:
                            j = skip_field(payload, j, de, kw)
                    if ids:
                        cols = []
                        for arr in (ids, lats, lons):
                            dec = []; pid = 0
                            for v in arr:
                                pid += v; dec.append(pid)
                            cols.append(dec)
                        group['dense'].extend(zip(cols[0], cols[1], cols[2]))
                elif gn == 3 and gw == 2:                          # Way
                    wl, j = read_varint(payload, j); we = j + wl
                    wid = 0; refs = keys = vals = None
                    while j < we:
                        k, j = read_varint(payload, j); kn = k >> 3; kw = k & 7
                        if kn == 1 and kw == 0:
                            wid, j = read_varint(payload, j)
                        elif kn == 2 and kw == 2:
                            kl, j = read_varint(payload, j); keys, j = packed_uint(payload, j, j + kl)
                        elif kn == 3 and kw == 2:
                            kl, j = read_varint(payload, j); vals, j = packed_uint(payload, j, j + kl)
                        elif kn == 8 and kw == 2:
                            kl, j = read_varint(payload, j); r, j = packed_sint(payload, j, j + kl)
                            dec = []; pid = 0
                            for v in r:
                                pid += v; dec.append(pid)
                            refs = dec
                        else:
                            j = skip_field(payload, j, we, kw)
                    if refs:
                        group['ways'].append((wid, refs, keys or [], vals or []))
                else:
                    j = skip_field(payload, j, end, gw)
            i = end
            out.append((group, strings, gran, lato, lono))
        else:
            i = skip_field(payload, i, n, wt)
    return out

# ---------------- file-level ----------------
def extract_blob_payload(blob):
    """Blob { raw=1; raw_size=2; zlib_data=3; zstd_data=7 } -> message bytes"""
    i = 0; n = len(blob)
    while i < n:
        f, i = read_varint(blob, i); num = f >> 3; wt = f & 7
        if num == 1 and wt == 2:
            l, i = read_varint(blob, i)
            return blob[i:i+l]
        elif num == 3 and wt == 2:
            l, i = read_varint(blob, i)
            return zlib.decompress(blob[i:i+l])
        elif num == 7 and wt == 2:
            l, i = read_varint(blob, i)
            try:
                import zstandard
                return zstandard.ZstdDecompressor().decompress(blob[i:i+l])
            except Exception:
                return None
        else:
            i = skip_field(blob, i, n, wt)
    return None

def iter_blocks(data):
    """yield (group, strings, gran, lato, lono) for each OSMData group"""
    off = 0
    while off < len(data):
        (hlen,) = struct.unpack('>I', data[off:off+4]); off += 4
        i = off; end = off + hlen
        btype = ''; dsz = 0
        while i < end:
            f, i = read_varint(data, i); num = f >> 3; wt = f & 7
            if num == 1 and wt == 2:
                l, i = read_varint(data, i); btype = data[i:i+l].decode(); i += l
            elif num == 3 and wt == 0:
                dsz, i = read_varint(data, i)
            else:
                i = skip_field(data, i, end, wt)
        off += hlen
        blob = data[off:off+dsz]; off += dsz
        if btype != 'OSMData':
            continue
        payload = extract_blob_payload(blob)
        if payload is None:
            continue
        for group, strings, gran, lato, lono in parse_block(payload):
            yield group, strings, gran, lato, lono

# ---------------- main ----------------
def main():
    src, dst = sys.argv[1], sys.argv[2]
    boxes = [tuple(float(x) for x in a.split(',')) for a in sys.argv[3:]]
    data = open(src, 'rb').read()
    print(f'file {len(data)/1e6:.0f} MB, {len(boxes)} box(es)', file=sys.stderr)

    def in_any(lat, lon):
        return any(minlo <= lon <= maxlo and minla <= lat <= maxla
                   for minlo, minla, maxlo, maxla in boxes)

    # pass 1: node ids inside the boxes
    print('pass 1: scanning nodes', file=sys.stderr)
    inside = set()
    for group, strings, gran, lato, lono in iter_blocks(data):
        conv = gran / 1e9
        for nid, la, lo in group['nodes']:
            if in_any((la + lato) * conv, (lo + lono) * conv):
                inside.add(nid)
        for nid, la, lo in group['dense']:
            if in_any((la + lato) * conv, (lo + lono) * conv):
                inside.add(nid)
    print(f'  {len(inside)} nodes in box(es)', file=sys.stderr)

    # pass 2: ways touching the boxes (tags decoded to strings now —
    # string indices are only valid inside their own block)
    print('pass 2: scanning ways', file=sys.stderr)
    ways = []
    referenced = set()
    for group, strings, gran, lato, lono in iter_blocks(data):
        for wid, refs, keys, vals in group['ways']:
            if any(r in inside for r in refs):
                tags = tuple((strings[k], strings[v]) for k, v in zip(keys, vals))
                ways.append((wid, refs, tags))
                referenced.update(refs)
    print(f'  {len(ways)} ways kept, {len(referenced)} referenced nodes', file=sys.stderr)

    # pass 3: positions of the kept nodes (raw PBF units)
    print('pass 3: collecting kept node positions', file=sys.stderr)
    kept = {}
    gran, lato, lono = 100, 0, 0
    for group, strings, gran_b, lato_b, lono_b in iter_blocks(data):
        gran, lato, lono = gran_b, lato_b, lono_b
        for nid, la, lo in group['nodes']:
            if nid in referenced: kept[nid] = (la, lo)
        for nid, la, lo in group['dense']:
            if nid in referenced: kept[nid] = (la, lo)
    print(f'  {len(kept)} positions (gran {gran}, offset {lato}/{lono})', file=sys.stderr)
    if len(kept) != len(referenced):
        print(f'  warning: {len(referenced) - len(kept)} referenced nodes not found', file=sys.stderr)

    # build a fresh string table from the strings actually used
    used = set()
    for _, _, tags in ways:
        for k, v in tags:
            used.add(k); used.add(v)
    strings_out = [''] + sorted(used)
    stridx = {t: i for i, t in enumerate(strings_out)}
    print(f'  string table: {len(strings_out)} strings', file=sys.stderr)

    # ---- write output ----
    out = bytearray()

    def fileblock(btype, payload):
        nonlocal out
        body = f_varint(2, len(payload)) + f_msg(3, zlib.compress(payload))
        hdr = f_msg(1, btype.encode()) + f_varint(3, len(body))
        out += struct.pack('>I', len(hdr)) + hdr + body

    # header block with combined bbox (nanodegrees)
    minlo = min(b[0] for b in boxes); minla = min(b[1] for b in boxes)
    maxlo = max(b[2] for b in boxes); maxla = max(b[3] for b in boxes)
    hdr = f_msg(1, f_sint(1, int(minlo * 1e9)) + f_sint(2, int(maxlo * 1e9))
                + f_sint(3, int(maxla * 1e9)) + f_sint(4, int(minla * 1e9)))
    fileblock('OSMHeader', hdr)

    def string_table_msg():
        return f_msg(1, b''.join(f_msg(1, s.encode()) for s in strings_out))

    # nodes as dense blocks (columns delta-coded)
    ids_sorted = sorted(kept)
    BLOCK = 2_000_000
    for b0 in range(0, len(ids_sorted), BLOCK):
        chunk = ids_sorted[b0:b0+BLOCK]
        d_ids = []; pid = 0
        for nid in chunk:
            d_ids.append(nid - pid); pid = nid
        d_la = []; pid = 0
        for nid in chunk:
            la, lo = kept[nid]; d_la.append(la - pid); pid = la
        d_lo = []; pid = 0
        for nid in chunk:
            la, lo = kept[nid]; d_lo.append(lo - pid); pid = lo
        dense = f_packed_sint(1, d_ids) + f_packed_sint(8, d_la) + f_packed_sint(9, d_lo)
        pb = string_table_msg() + f_msg(2, f_msg(2, dense)) + f_varint(17, 100)
        fileblock('OSMData', pb)

    # ways blocks: chunk of 50k ways per PrimitiveGroup
    CHUNK = 50_000
    def way_msg(wid, refs, tags):
        d = []; pid = 0
        for r in refs:
            d.append(r - pid); pid = r
        w = f_varint(1, wid)
        if tags:
            w += f_packed_uint(2, [stridx[k] for k, v in tags])
            w += f_packed_uint(3, [stridx[v] for k, v in tags])
        w += f_packed_sint(8, d)
        return w

    for w0 in range(0, len(ways), CHUNK):
        chunk = ways[w0:w0+CHUNK]
        body = string_table_msg()
        body += f_msg(2, b''.join(f_msg(3, way_msg(*w)) for w in chunk))
        body += f_varint(17, 100)
        fileblock('OSMData', body)

    open(dst, 'wb').write(bytes(out))
    print(f'wrote {dst}: {len(ids_sorted)} nodes, {len(ways)} ways, {len(out)/1e6:.1f} MB', file=sys.stderr)

if __name__ == '__main__':
    main()