import json,gzip,subprocess,hashlib,time,os,collections
from pathlib import Path
from mysql_encoding import PROBE_SQL,validate_encoding,consume_probe
from data_contract import canonical_row,row_digest,quote_identifier,select_expression
r=Path(__file__).resolve().parents[1]
b=json.loads((r/'private/mysql-original-baseline.json').read_text())
validate_encoding(b.get('encoding'))
folder=r/'private/original-export';folder.mkdir(mode=0o700,exist_ok=True)
cols=collections.defaultdict(list)
for c in b['columns']:cols[c['TABLE_NAME']].append(c)
parts=[PROBE_SQL,"SET SESSION time_zone='+08:00'; START TRANSACTION WITH CONSISTENT SNAPSHOT;"]
for t in b['tables']:
    name=t['TABLE_NAME'];expr=','.join(select_expression(c)for c in cols[name])
    parts.append("SELECT JSON_OBJECT('table','"+name+"','row',JSON_ARRAY("+expr+")) FROM "+quote_identifier(name)+';')
parts.append('COMMIT;')
query='\n'.join(parts).encode('utf8')
cmd=['docker','exec','-i','foodsafe2-mysql8-stage-20261023','mysql','--defaults-extra-file=/run/secrets/client.cnf','--default-character-set=utf8mb4','--database=foodsafe2','--batch','--raw','--quick','--skip-column-names','--max-allowed-packet=1G']
log=r/'private/source-export.stderr';log_handle=log.open('wb');os.chmod(log,0o600)
p=subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=log_handle)
p.stdin.write(query);p.stdin.close()
encoding=consume_probe(p)
counts={k:0 for k in b['counts']};hashes=collections.defaultdict(list);files={};max_lengths=collections.defaultdict(dict);t0=time.time();last=t0;total=0
try:
    for line in p.stdout:
        obj=json.loads(line);name=obj['table'];row=obj['row']
        if len(row)!=len(cols[name]):raise RuntimeError('Column count mismatch')
        if name not in files:
            files[name]=gzip.open(folder/(name+'.jsonl.gz.part'),'wb',compresslevel=1)
            os.chmod(folder/(name+'.jsonl.gz.part'),0o600)
        raw=canonical_row(row);files[name].write(raw+b'\n');counts[name]+=1;total+=1;hashes[name].append(hashlib.sha256(raw).digest())
        for c,value in zip(cols[name],row):
            if value is not None:
                size=len(value.encode('utf8'));key=c['COLUMN_NAME'];max_lengths[name][key]=max(max_lengths[name].get(key,0),size)
        if time.time()-last>=20:
            print(json.dumps({'phase':'export_original','table':name,'rows_exported':total,'elapsed_s':round(time.time()-t0)}),flush=True);last=time.time()
    rc=p.wait()
    if rc:raise RuntimeError('Source export failed; see private stderr')
    mismatch={k:{'expected':v,'actual':counts[k]}for k,v in b['counts'].items()if int(v)!=counts[k]}
    if mismatch:raise RuntimeError('Source export row-count mismatch: '+json.dumps(mismatch))
finally:
    for f in files.values():f.close()
    log_handle.close()
summary={'encoding':encoding,'table_count':len(counts),'total_rows':total,'counts':counts,'max_utf8_lengths':dict(max_lengths),'multiset_sha256':{k:hashlib.sha256(b''.join(sorted(hashes[k]))).hexdigest()for k in counts},'elapsed_s':round(time.time()-t0),'exported_as_strings_without_js_numeric_conversion':True,'mysql_source_read_only':True}
for name in files:(folder/(name+'.jsonl.gz.part')).rename(folder/(name+'.jsonl.gz'))
(r/'private/original-export-manifest.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2))
print(json.dumps({'export_complete':True,'total_rows':total,'tables':len(counts),'elapsed_s':summary['elapsed_s']}),flush=True)
