"""Re-fetch a bounded public Gaia DR3/2MASS study. No account or API key.
Run with Python3; source request bytes and query URLs are retained for provenance.
"""
from pathlib import Path
from urllib.request import Request,urlopen
from urllib.parse import urlencode
import json,csv,io,hashlib,datetime,time
BASE=Path(__file__).resolve().parent
ENDPOINT='https://gea.esac.esa.int/tap-server/tap/sync'
COLUMNS=['source_id','ra','dec','ref_epoch','ra_error','dec_error','parallax','parallax_error','pmra','pmra_error','pmdec','pmdec_error','ra_dec_corr','ra_parallax_corr','dec_parallax_corr','parallax_pmra_corr','parallax_pmdec_corr','ra_pmra_corr','ra_pmdec_corr','dec_pmra_corr','dec_pmdec_corr','pmra_pmdec_corr','phot_g_mean_mag','bp_rp','radial_velocity','radial_velocity_error','ruwe']
QUERY='SELECT TOP 20 '+','.join(COLUMNS)+' FROM gaiadr3.gaia_source WHERE parallax_over_error>20 AND ruwe<1.4 AND phot_g_mean_mag BETWEEN 8 AND 13 AND pm>300 AND dec BETWEEN -60 AND 60 ORDER BY random_index'
def read(name,query):
 url=ENDPOINT+'?'+urlencode({'REQUEST':'doQuery','LANG':'ADQL','FORMAT':'csv','QUERY':query})
 with urlopen(Request(url,headers={'User-Agent':'SpaceTRC-Research/9.0 public archive study','Accept':'text/csv'}),timeout=60) as r:
  raw=r.read(2_000_001)
  if r.status!=200 or len(raw)>2_000_000:raise ValueError('Bounded TAP response failed')
 text=raw.decode('utf-8-sig')
 if 'QUERY_STATUS' in text or text.lstrip().startswith('<'):raise ValueError(text[:600])
 (BASE/name).write_bytes(raw)
 rows=list(csv.DictReader(io.StringIO(text)))
 return rows,{'file':name,'url':url,'adql':query,'checkedAt':datetime.datetime.now(datetime.timezone.utc).isoformat(),'bytes':len(raw),'sha256':hashlib.sha256(raw).hexdigest(),'rows':len(rows),'provider':'ESA Gaia Archive TAP'}
if __name__=='__main__':
 selection=list(csv.DictReader(io.StringIO((BASE/'gaia-dr3-selection.csv').read_text())))
 selection_ids=','.join(r['source_id'] for r in selection)
 query='SELECT '+','.join(COLUMNS)+' FROM gaiadr3.gaia_source WHERE source_id IN ('+selection_ids+') ORDER BY source_id'
 rows,source=read('gaia-dr3-source.csv',query)
 (BASE/'gaia-source-provenance.json').write_text(json.dumps(source,ensure_ascii=False,indent=2)+'\n')
 assert len(rows)==20
 ids=','.join(r['source_id'] for r in rows)
 schema=list(csv.DictReader(io.StringIO((BASE/'tmass-schema.csv').read_text())))
 metadata=json.loads((BASE/'tmass-schema-provenance.json').read_text())
 tmq="SELECT TOP 100 n.source_id,n.original_ext_source_id,n.angular_distance,n.score,n.xm_flag,n.clean_tmass_psc_xsc_oid,t.designation,t.ra AS tm_ra,t.dec AS tm_dec,t.j_date,t.err_maj,t.err_min,t.err_ang,t.j_m,t.j_msigcom,t.h_m,t.h_msigcom,t.ks_m,t.ks_msigcom,t.ph_qual FROM gaiadr3.tmass_psc_xsc_neighbourhood AS n JOIN gaiadr3.tmass_psc_xsc_join AS j USING(clean_tmass_psc_xsc_oid) JOIN gaiadr1.tmass_original_valid AS t ON j.original_psc_source_id=t.designation WHERE n.source_id IN ("+ids+")"
 tmass,tmprov=read('tmass-neighbours.csv',tmq)
 (BASE/'tmass-provenance.json').write_text(json.dumps(tmprov,ensure_ascii=False,indent=2)+'\n')
 (BASE/'gaia-source-provenance.json').write_text(json.dumps(source,ensure_ascii=False,indent=2)+'\n')
 (BASE/'tmass-schema-provenance.json').write_text(json.dumps(metadata,ensure_ascii=False,indent=2)+'\n')
 print(json.dumps({'source':source,'tmassColumns':schema},ensure_ascii=False),flush=True)