#!/usr/bin/env python3
"""Frozen-anchor, trace-local brackets for the selected detour review cohort.

No inferred target is promoted to an anchor. Compact evidence contains one
observation per (trace, target, mode, ordered city pair), not one per probe.
"""
import collections as co
import concurrent.futures as cf
import functools
import gzip
import hashlib
import html
import ipaddress
import json
import math
import os
from pathlib import Path
import subprocess
import sys
import time
import orjson as J

sys.dont_write_bytecode = True
HERE = Path(__file__).resolve().parent
PAPER = HERE.parent
ROOT = PAPER.parent
RAW = PAPER / 'as14593_internal_first_round_20261108/partitions'
PAGE = PAPER / 'early-late-ingress/early_review_rounds/data.json'
KEEP = ['支持早入', '第一跳未定位', 'ic 无位置']
MODES = ['high_medium', 'high_only']
ALIASES = {'São Paulo/SP':'São Paulo','Frankfurt am Main':'Frankfurt',
           'Montréal':'Montreal','Minneapolis, MN':'Minneapolis','Bogota':'Bogotá',
           'Queretaro':'Querétaro','Brasília/DF':'Brasília','Fortaleza/CE':'Fortaleza'}
NONCITY = {'Amsterdam, Rotterdam, Brussels, Luxembourg, Frankfurt, Düsseldorf, Berlin, Dublin, London, Copenhagen, Paris, Marseille',
           'Los Angeles and Silicon Valley','Northern Virginia','Argentina','West Coast'}

def read(p):
    with (gzip.open(p,'rb') if str(p).endswith('.gz') else open(p,'rb')) as f:
        return J.loads(f.read())

def write(p,obj,pretty=False):
    b=J.dumps(obj,option=J.OPT_INDENT_2 if pretty else 0)
    if str(p).endswith('.gz'):
        with gzip.open(p,'wb',compresslevel=1) as f: f.write(b)
    else: Path(p).write_bytes(b)

def sha(p): return hashlib.sha256(Path(p).read_bytes()).hexdigest()
def canon(s): return ALIASES.get(s,s)

@functools.lru_cache(None)
def is_public(ip):
    try: return ipaddress.ip_address(ip).is_global
    except ValueError: return False

def prep():
    data=read(PAGE)
    rows=[t for t in data['traces'] if t['ptype']=='绕' and t['cat'] in KEEP]
    multi={h[1] for t in rows for h in t['internal'] if h[4]=='多城市'}
    targets={h[1] for t in rows for h in t['internal'] if h[4]!='已定位'}-multi
    assert (len(rows),len(multi),len(targets)) == (7204,79,910)
    sys.path.insert(0,str(PAPER/'pop_bracket_ttl5_14_20261111'))
    import analyze_brackets as A
    sys.path.insert(0,str(PAPER/'city_group_rebuild_20261108'))
    import prepare_current as C
    assert C.rdns_check.USE_CONFIRMED and Path(C.rdns_check.CONFIRMED_CSV).exists()
    blocked=set(multi)|targets
    # Exclude every multi-city IP in the rounds, not just those in this cohort.
    for book in (A.a,A.b):
        for ip,r in book.items():
            if r.get('城市数',r.get('城市名称数',0))>1: blocked.add(ip)
    anchors={}; excluded=[];coords={}
    for ip,a in (A.ANCHORS|A.IXP).items():
        a=dict(a,original_city=a['city'],city=canon(a['city']))
        if ip in blocked or a['original_city'] in NONCITY:
            excluded.append(dict(a,ip=ip,reason='target_or_multicity' if ip in blocked else 'not_a_unique_city'));continue
        ck=a['city']+'|'+a['country']
        ll=A.COORD.get(a['city']) or C.rdns_check.latlon(a['city'],a['country'])
        if ll is None:
            excluded.append(dict(a,ip=ip,reason='no_unambiguous_city_coordinates'));continue
        coords[ck]=list(ll)
        anchors[ip]=dict(a,city_key=ck,coordinate=list(ll))
    provenance=co.defaultdict(list)
    for g in read(PAPER/'round1_round2_strict_rebuild_20261109/round1_groups.json.gz'):
        ip=g['ip']
        if ip in anchors and canon(g['city'])==anchors[ip]['city'] and g['confidence'] in ('high','medium'):
            s=g['sample']
            provenance[ip].append({'round':1,'confidence':g['confidence'],'trace_count':g['trace_count'],
               'days':g['days'],'representative_ic_ip':s[7],'representative_ic_city_source':s[9],
               'representative_endpoint_source':s[12],'representative_trace_id':s[-1]})
    for ip,a in anchors.items():
        a['upstream_round1_representatives']=provenance.get(ip,[])
        a['source_note']='冻结轮次定位；上游可能含 IPinfo，代表证据来源不等于完整溯源审计' if a['kind'].startswith('round') else 'PeeringDB Starlink 接口；登记城市作为独立锚点'
    targets=sorted(targets,key=lambda x:(ipaddress.ip_address(x).version,int(ipaddress.ip_address(x))))
    inventory={'anchors':anchors,'excluded':excluded,'coordinates':coords,'aliases':ALIASES,
      'rdns_policy':'确认库 > rDNS脚本（待复核） > IPinfo（无/未识别 rDNS 时）。原冻结轮次锚点置信度与本次 rDNS 参考来源分开记录；IPinfo/脚本参考不会自动成为中高置信锚点。',
      'confirmed_csv':str(C.rdns_check.CONFIRMED_CSV),'confirmed_sha256':sha(C.rdns_check.CONFIRMED_CSV)}
    write(HERE/'anchor_inventory.json',inventory,True)
    write(HERE/'selected_traces.json.gz',rows)
    target_rows=[]
    for ip in targets:
        hit=[t for t in rows if any(h[1]==ip for h in t['internal'])]
        target_rows.append({'ip':ip,'selected_trace_count':len(hit),'review_categories':dict(co.Counter(t['cat'] for t in hit)),
            'primary_labels':dict(co.Counter(t['label'] for t in hit)),'selected_trace_ids':[t['id'] for t in hit]})
    manifest={'page':str(PAGE.with_name('early_review.html')),'data':str(PAGE),'data_sha256':sha(PAGE),
      'categories':KEEP,'selected_trace_count':len(rows),'selected_group_count':len({t['gid'] for t in rows}),
      'category_counts':dict(co.Counter(t['cat'] for t in rows)),'primary_label_counts':dict(co.Counter(t['label'] for t in rows)),
      'selected_traces_with_targets':sum(any(h[1] in targets for h in t['internal']) for t in rows),
      'target_count':len(targets),'targets':targets,'target_rows':target_rows,'excluded_multicity_ips':sorted(multi),
      'anchor_count':len(anchors),'anchor_kind_counts':dict(co.Counter(a['kind'] for a in anchors.values())),
      'raw_partitions':[str(p) for p in sorted(RAW.glob('????-??-??.jsonl.gz'))],
      'raw_trace_count':sum(read(p.with_suffix('.report.json'))['stats']['matching_traces'] for p in RAW.glob('????-??-??.jsonl.gz'))}
    write(HERE/'scope_manifest.json',manifest,True)
    (HERE/'patterns.txt').write_text('\n'.join(ip+'"' for ip in targets)+'\n'+'\n'.join(t['id']+'"' for t in rows)+'\n')
    print(J.dumps({k:v for k,v in manifest.items() if k not in ('targets','target_rows','raw_partitions','excluded_multicity_ips')}).decode(),flush=True)

def init():
    global INV,ANCHORS,ANCHOR_IPS,AI,TARGETS,TI,TARGET_SET,SELECTED,CITY_KEYS,CI,CDIST,BASESTAR
    INV=read(HERE/'anchor_inventory.json'); ANCHORS=INV['anchors'];ANCHOR_IPS=sorted(ANCHORS);AI={ip:i for i,ip in enumerate(ANCHOR_IPS)}
    m=read(HERE/'scope_manifest.json');TARGETS=m['targets'];TI={ip:i for i,ip in enumerate(TARGETS)};TARGET_SET=set(TARGETS)
    SELECTED={r['id']:r for r in read(HERE/'selected_traces.json.gz')}
    CITY_KEYS=sorted(INV['coordinates']);CI={c:i for i,c in enumerate(CITY_KEYS)};CDIST={}
    for a in CITY_KEYS:
        for b in CITY_KEYS:
            la,lo=map(math.radians,INV['coordinates'][a]);lb,lob=map(math.radians,INV['coordinates'][b])
            v=math.sin((lb-la)/2)**2+math.cos(la)*math.cos(lb)*math.sin((lob-lo)/2)**2
            CDIST[(CI[a],CI[b])]=12742*math.asin(min(1,math.sqrt(v)))
    BASESTAR=set(ANCHORS)|TARGET_SET

def path_grade(r,hs):
    if r.get('source_kind')!='cloud': return 'high',None
    universe=BASESTAR|set(r.get('candidates',[]))|set(r.get('unknown_candidates',[]));meta=r.get('meta',{})
    star=lambda ip:ip in universe or str(meta.get(ip,{}).get('asn','')).upper()=='AS14593'
    st=[t for t,ips in hs.items() if any(star(ip) for ip in ips)]
    if not st:return 'invalid',0
    egress=r.get('direction')=='egress';boundary=max(st) if egress else min(st)
    n=sum(any(is_public(ip) and not star(ip) for ip in ips) for t,ips in hs.items() if (t>boundary if egress else t<boundary))
    return ('high' if n>=4 else 'medium' if n>=2 else 'invalid'),n

def obs_rank(o):
    # obs = target index, mode, left city index, right city index, target TTL,
    # left anchor index, left TTL, right anchor index, right TTL.
    return (sum(ANCHORS[ANCHOR_IPS[i]]['confidence']!='high' for i in (o[5],o[7])),o[8]-o[6],o[4],o[5],o[7],o[6],o[8])

def observations(r):
    hs=co.defaultdict(set)
    for h in r['hops']:
        ip=h.get('ip');t=h.get('ttl')
        if ip not in (None,'*') and isinstance(t,int) and 0<t<255 and not h.get('ttl_uncertain'):hs[t].add(ip)
    found=TARGET_SET&{ip for ips in hs.values() for ip in ips}
    grade,outside=path_grade(r,hs)
    if not found or grade=='invalid' or r.get('same_id_content_conflict'):return {},hs,found,grade,outside
    loc=[(t,ip,AI[ip],CI[ANCHORS[ip]['city_key']]) for t,ips in sorted(hs.items()) if len(ips)==1 for ip in ips if ip in ANCHORS and ip not in (r.get('src'),r.get('dst'))]
    # Interval validity is shared by every target between the same two anchors.
    intervals=[]
    for i,(lt,li,la,lc) in enumerate(loc):
        for rt,ri,ra,rc in loc[i+1:]:
            if any(len(hs.get(t,set()))>1 for t in range(lt,rt+1)):continue
            if sum(t not in hs for t in range(lt+1,rt))>2:continue
            modes=(0,1) if ANCHORS[li]['confidence']==ANCHORS[ri]['confidence']=='high' else (0,)
            intervals.append((lt,la,lc,rt,ra,rc,modes))
    local={}
    for tt,ips in hs.items():
        if len(ips)!=1:continue
        for ip in ips&TARGET_SET:
            if ip in (r.get('src'),r.get('dst')):continue
            ti=TI[ip]
            for lt,la,lc,rt,ra,rc,modes in intervals:
                if lt<tt<rt:
                    for mode in modes:
                        o=[ti,mode,lc,rc,tt,la,lt,ra,rt];key=(ti,mode,lc,rc)
                        if key not in local or obs_rank(o)<obs_rank(local[key]):local[key]=o
    return local,hs,found,grade,outside

def scan_day(path):
    init();p=Path(path);day=p.name[:10];outdir=HERE/'daily';outdir.mkdir(exist_ok=True)
    done=outdir/(day+'.summary.json')
    if done.exists():return read(done)
    start=time.monotonic();stats=co.Counter();per_ip=co.defaultdict(co.Counter);groups={};seen=set();metadata={}
    selected_found=[];pop_cache={};pop_feeds={}
    sys.path.insert(0,str(PAPER/'city_group_rebuild_20261108'))
    import prepare_current as C
    def pop_lookup(r):
        stamp=r.get('start_utc') or '';utcday=stamp[:10]
        if len(stamp)!=20:return {'status':'missing_measurement_time'}
        feed=pop_feeds.setdefault(utcday,None)
        if feed is None:feed=pop_feeds[utcday]=C.DailyPops(utcday)
        # Hour granularity is sufficient: all snapshots fall exactly on an hour.
        key=(r.get('dst',''),stamp[:13])
        if key not in pop_cache:
            z=feed.lookup(key[0],stamp);z['source']='daily_public_pop';z['measurement_utc_day']=utcday
            pop_cache[key]=z
        return pop_cache[key]
    def meta_note(r,ip):
        m=r.get('meta',{}).get(ip,{})
        v=[ip,m.get('rdns') or '',str(m.get('asn') or ''),m.get('ipinfo_city',m.get('city')) or '']
        k=tuple(v)
        if k in metadata:metadata[k]['trace_count']+=1
        else:metadata[k]={'ip':ip,'rdns':v[1],'asn':v[2],'ipinfo_city':v[3],'trace_count':1,'example_trace_id':r['trace_id'],'day':day}
    cmd=['rg','--search-zip','--line-number','--no-filename','--fixed-strings','-f',str(HERE/'patterns.txt'),str(p)]
    proc=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    with gzip.open(outdir/(day+'.evidence.jsonl.gz'),'wb',compresslevel=1) as out, gzip.open(outdir/(day+'.selected_raw.jsonl.gz'),'wb',compresslevel=1) as selected_out:
        for ln in proc.stdout:
            lineno,line=ln.split(b':',1);r=J.loads(line);stats['text_matches']+=1
            found=TARGET_SET&{h.get('ip') for h in r.get('hops',[])};selected=r['trace_id'] in SELECTED
            if not found and not selected:continue
            rk=(r['trace_id'],r.get('fingerprint',''))
            if rk in seen:stats['duplicate_records']+=1;continue
            seen.add(rk);stats['matching_traces']+=1
            if selected:
                r['_partition_file']=str(p);r['_partition_line']=int(lineno)
                selected_out.write(J.dumps(r)+b'\n');selected_found.append(r['trace_id'])
            local,hs,valid_found,grade,outside=observations(r)
            for ip in found:
                ps=per_ip[ip];ps['traces']+=1;ps['path_'+grade]+=1
                if selected:ps['selected_raw_traces']+=1
                if ip not in valid_found:ps['only_uncertain_ttl']+=1
                if ip in (r.get('src'),r.get('dst')):ps['endpoint_occurrence']+=1
                if r.get('same_id_content_conflict'):ps['content_conflict']+=1
            good={TARGETS[k[0]] for k in local}
            for ip in good:per_ip[ip]['eligible_traces']+=1
            if local:
                stats['eligible_traces']+=1;stats['observations']+=len(local)
                ctx={'trace_id':r['trace_id'],'day':day,'start_utc':r.get('start_utc'),'src':r.get('src'),'raw_src':r.get('raw_src'),
                  'node':r.get('node'),'source_kind':r.get('source_kind'),'direction':r.get('direction'),
                  'dst':r.get('dst'),'destination_pop':pop_lookup(r),'path_grade':grade,'outside_public_ttls':outside,
                  'selected_cohort':selected,'partition_file':str(p),'partition_line':int(lineno),
                  'source_file':r.get('source_file'),'source_index':r.get('source_index'),
                  'repeated_target_ips':[ip for ip in good if sum(ip in ips for ips in hs.values())>1]}
                if selected:
                    s=SELECTED[r['trace_id']];ctx['review_category']=s['cat'];ctx['original_label']=s['label'];ctx['page_pop']=s['pop'];ctx['page_pop_city']=s['pop_city']
                out.write(J.dumps({'context':ctx,'observations':list(local.values())})+b'\n')
                used=set(found)
                for o in local.values():used.update((ANCHOR_IPS[o[5]],ANCHOR_IPS[o[7]]))
                for ip in used:meta_note(r,ip)
                for key,o in local.items():
                    if key not in groups:groups[key]={'n':0,'rank':None,'sample':None,'selected_n':0,'conf':co.Counter(),'grades':co.Counter()}
                    g=groups[key];g['n']+=1;g['selected_n']+=selected
                    conf='/'.join(ANCHORS[ANCHOR_IPS[i]]['confidence'] for i in (o[5],o[7]));g['conf'][conf]+=1;g['grades'][grade]+=1
                    rank=(*obs_rank(o)[:2],not selected,grade!='high',r['trace_id'],o[5],o[7],o[4])
                    if g['rank'] is None or rank<g['rank']:g['rank']=rank;g['sample']={'context':ctx,'observation':o}
            elif found:
                for ip in found:meta_note(r,ip)
            if stats['matching_traces']%100000==0:
                print(J.dumps({'day':day,'processed':stats['matching_traces'],'eligible':stats['eligible_traces'],'seconds':round(time.monotonic()-start)}).decode(),flush=True)
    error=proc.stderr.read().decode();code=proc.wait()
    if code not in (0,1):raise RuntimeError((day,code,error))
    write(outdir/(day+'.groups.json.gz'),[dict(key=list(k),**v) for k,v in groups.items()])
    write(outdir/(day+'.metadata.json.gz'),list(metadata.values()))
    result={'day':day,'stats':dict(stats),'per_ip':dict(per_ip),'selected_found':selected_found,
      'groups':len(groups),'raw_traces':read(p.with_suffix('.report.json'))['stats']['matching_traces'],
      'seconds':round(time.monotonic()-start,1),'source_size':p.stat().st_size}
    write(done,result,True);print(J.dumps({k:v for k,v in result.items() if k not in ('per_ip','selected_found')}).decode(),flush=True)
    return result

def scan():
    files=read(HERE/'scope_manifest.json')['raw_partitions']
    if len(sys.argv)>2:files=[p for p in files if Path(p).name[:10] in sys.argv[2:]]
    files.sort(key=lambda p:(Path(p).name[:10]!='2026-10-31',p))
    with cf.ProcessPoolExecutor(max_workers=int(os.environ.get('BRACKET_WORKERS','4'))) as pool:
        results=list(pool.map(scan_day,files))
    write(HERE/'scan_manifest.json',results,True)

if __name__=='__main__':
    {'prep':prep,'scan':scan}[sys.argv[1]]()