"""Full polygons, z10 with adaptive subdivision; bounded, deterministic gzip JSON tiles."""
import sys,pathlib,json,math,gzip,hashlib,shutil,collections,threading,concurrent.futures
import pyarrow.parquet as pq
import orjson
from shapely import from_wkb
from export import ROOT,REGIONS
RELEASE='2026-09-23.1'; REV='v1-20260923'
OUT=ROOT/'data'/'site'; RAW=ROOT/'data'/'buckets'
MAX_RAW=2_000_000

def xy(lon,lat,z):
 n=2**z; return int((lon+180)/360*n),int((1-math.asinh(math.tan(math.radians(lat)))/math.pi)/2*n)
def bounds(z,x,y):
 n=2**z
 return [x/n*360-180,math.degrees(math.atan(math.sinh(math.pi*(1-2*(y+1)/n)))),(x+1)/n*360-180,math.degrees(math.atan(math.sinh(math.pi*(1-2*y/n))))]
def keys(b,z):
 x0,y1=xy(b[0],b[1],z);x1,y0=xy(b[2],b[3],z)
 return [(z,x,y) for x in range(x0,x1+1) for y in range(y0,y1+1)]
def overlaps(a,b):return a[0]<=b[2] and a[2]>=b[0] and a[1]<=b[3] and a[3]>=b[1]
PREFIX=('{'+'"schema":1,"release":"'+RELEASE+'","buildings":[').encode()
def encoded_size(rows):return len(PREFIX)+2+sum(len(r['_raw']) for r in rows)+max(0,len(rows)-1)
def encode(rows):return PREFIX+b','.join(r['_raw'] for r in rows)+b']}'
def load_rows(path):
 rows=[]
 if path.exists():
  with path.open('rb') as f:
   for line in f:
    raw=line.rstrip(b'\n');r=orjson.loads(raw) # strict JSON: rejects non-finite numbers
    rows.append({'id':r['id'],'b':r['b'],'_raw':raw})
 return rows

def main():
 OUT.mkdir(parents=True,exist_ok=True); RAW.mkdir(parents=True,exist_ok=True)
 extracts=sorted((ROOT/'data'/'extracts').glob('*.parquet')); regions={f.stem:REGIONS[f.stem] for f in extracts}
 stage=RAW/'stage.json'
 staged=json.loads(stage.read_text()) if stage.exists() else {}
 sources=set(); count=0; handles=collections.OrderedDict()
 # Reuse completed regional staging. Source Parquet files are immutable in this release.
 for f in extracts:
  if f.stem in staged:
   old=staged[f.stem];assert old['bytes']==f.stat().st_size,'Changed source: rebuild staging'
   count+=old['count'];sources.update(old['sources']);continue
  before=count;region_sources=set()
  for batch in pq.ParquetFile(f).iter_batches(batch_size=10000):
   for row in batch.to_pylist():
    geom=from_wkb(row['geometry']);geom=geom.simplify(.000002,preserve_topology=True)
    parts=[geom] if geom.geom_type=='Polygon' else list(geom.geoms) if geom.geom_type=='MultiPolygon' else []
    ss=json.loads(row['sources']); datasets=sorted(set(s['dataset'] for s in ss));sources.update(datasets);region_sources.update(datasets)
    prop='height' if row['height'] is not None else 'num_floors'
    hs=[s for s in ss if s.get('property') in (prop,'/'+prop)] or [s for s in ss if s.get('property')=='']
    exact=bool(hs) and all(s.get('dataset')=='OpenStreetMap' for s in hs)
    for i,p in enumerate(parts):
     if p.is_empty or not p.is_valid: raise ValueError('Invalid source geometry '+row['id'])
     rings=[[[round(x*1e6),round(y*1e6)] for x,y,*_ in r.coords] for r in [p.exterior,*p.interiors]]
     obj={'id':row['id']+':'+str(i),'r':rings,'h':row['height'],'l':row['num_floors'],'m':row['min_height'],'q':exact,'p':row['has_parts'],'s':datasets}
     b=list(p.bounds);obj['b']=b
     line=json.dumps(obj,separators=(',',':'))+'\n'
     for z,x,y in keys(b,10):
      k=f'{z}-{x}-{y}';
      if k not in handles:
       if len(handles)>=128:handles.popitem(last=False)[1].close()
       handles[k]=open(RAW/(k+'.jsonl'),'a')
      handles.move_to_end(k);handles[k].write(line)
     count+=1
  for h in handles.values():h.flush()
  staged[f.stem]={'count':count-before,'bytes':f.stat().st_size,'sources':sorted(region_sources)}
  stage.write_text(json.dumps(staged))
  print(f.stem,'polygons so far',count,flush=True)
 for h in handles.values():h.close()
 if '--stage-only' in sys.argv:return
 # Replace only this locally generated revision before first publication.
 if (OUT/REV).exists():shutil.rmtree(OUT/REV)
 tiles={}; total=0;maxfile=0;unique=0; output_lock=threading.Lock()
 def write(key,rows):
  nonlocal total,maxfile,unique
  z,x,y=key;rows=sorted({r['id']:r for r in rows}.values(),key=lambda r:r['id'])
  if encoded_size(rows)>MAX_RAW:
   if z>=16:raise ValueError('Tile exceeds client limit')
   for cx in (x*2,x*2+1):
    for cy in (y*2,y*2+1):
     b=bounds(z+1,cx,cy);write((z+1,cx,cy),[r for r in rows if overlaps(r['b'],b)])
   return
  payload=encode(rows)
  data=gzip.compress(payload,compresslevel=6,mtime=0);name=f'{z}/{x}/{y}.json.gz';path=OUT/REV/name;path.parent.mkdir(parents=True,exist_ok=True);path.write_bytes(data)
  with output_lock:
   tiles[f'{z}/{x}/{y}']={'bytes':len(data),'rawBytes':len(payload),'count':len(rows),'sha256':hashlib.sha256(data).hexdigest()}
   total+=len(data);maxfile=max(maxfile,len(data));unique+=len(rows)
 allkeys=set(k for b in regions.values() for k in keys(b,10))
 def process(key):
  f=RAW/('%d-%d-%d.jsonl'%key);rows=load_rows(f)
  write(key,rows)
 with concurrent.futures.ThreadPoolExecutor(max_workers=6) as pool:
  for _ in pool.map(process,sorted(allkeys)):pass
 tiles=dict(sorted(tiles.items()))
 manifest={'schema':1,'revision':REV,'release':RELEASE,'date':'2026-09-23T00:00:00Z','license':'ODbL-1.0','attribution':'© OpenStreetMap contributors, Overture Maps Foundation','sources':sorted(sources),'regions':[{'name':n,'bbox':b} for n,b in regions.items()],'tiles':tiles}
 (OUT/'manifest.json').write_text(json.dumps(manifest,separators=(',',':')))
 (OUT/'_headers').write_text(f'/{REV}/*\n  Cache-Control: public, max-age=31536000, immutable\n  Content-Type: application/gzip\n  Access-Control-Allow-Origin: *\n/manifest.json\n  Cache-Control: public, max-age=3600\n  Access-Control-Allow-Origin: *\n')
 (OUT/'404.html').write_text('Building tile unavailable')
 files=[p for p in OUT.rglob('*') if p.is_file()];sizes=[p.stat().st_size for p in files]
 stats={'regions':regions,'sourcePolygons':count,'tileCopies':unique,'tiles':len(tiles),'files':len(files),'totalBytes':sum(sizes),'largestFileBytes':max(sizes),'maxTileBytes':maxfile,'maxRawTileBytes':max(t['rawBytes'] for t in tiles.values())}
 (ROOT/'evidence'/'sizes.json').write_text(json.dumps(stats,indent=2));print(json.dumps(stats,indent=2))
 if '--candidate' not in sys.argv and (sum(sizes)>9_000_000_000 or len(files)>19000 or max(sizes)>=25*1024*1024):raise ValueError('Safety margin exceeded; run select_coverage.py before any deployment')
if __name__=='__main__':main()