#!/usr/bin/env python3
"""Deterministic LAION-DISCO audio tranche ingestion."""
from __future__ import annotations
import argparse, array, hashlib, heapq, json, math, os, re, shutil, subprocess, sys, time, unicodedata, wave
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
ROOT=Path("/home/ubuntu/minimax-laion-corpus")
REVISION="6e7bf3758a77301e46a715af894fefd79bb1da53"
CANDIDATES=ROOT/"manifests/candidates.jsonl"; CANDIDATE_SUMMARY=ROOT/"manifests/candidates.summary.json"
KNOWN_HASHES=ROOT/"manifests/known_user22_hashes.json"; RESULTS=ROOT/"manifests/ingest-results.jsonl"
CANONICAL=ROOT/"canonical"; TMP=ROOT/"tmp"; YTDLP=ROOT/".venv/bin/yt-dlp"
FFMPEG=shutil.which("ffmpeg") or "/usr/bin/ffmpeg"; FFPROBE=shutil.which("ffprobe") or "/usr/bin/ffprobe"
ID_RE=re.compile(r"^[A-Za-z0-9_-]{11}$")
def sha256_file(path:Path,chunk:int=8<<20)->str:
 h=hashlib.sha256()
 with path.open("rb") as f:
  while b:=f.read(chunk):h.update(b)
 return h.hexdigest()
def norm(value:Any)->str:
 if value is None:return ""
 return " ".join(unicodedata.normalize("NFKC",str(value)).casefold().split())
def stable_rank(song_id:str)->str:return hashlib.sha256((REVISION+"\0"+song_id).encode()).hexdigest()
def artist_keys(row:dict[str,Any])->list[str]:
 ids=sorted({"id:"+norm(x) for x in (row.get("artist_ids") or []) if norm(x)})
 if ids:return ids
 names=sorted({"name:"+norm(x) for x in (row.get("artist_names") or []) if norm(x)})
 return names or ["song:"+row["song_id"]]
def album_key(row:dict[str,Any])->str:
 aid=norm(row.get("album_id"))
 if aid:return "id:"+aid
 name=norm(row.get("album_name")); artists=artist_keys(row)
 if name:return "name:"+name+"|artist:"+(artists[0] if artists else row["song_id"])
 return "song:"+row["song_id"]
def atomic_text(path:Path,text:str)->None:
 path.parent.mkdir(parents=True,exist_ok=True);tmp=path.with_suffix(path.suffix+".tmp");tmp.write_text(text,encoding="utf-8");os.replace(tmp,path)
def write_json(path:Path,obj:Any)->None:atomic_text(path,json.dumps(obj,sort_keys=True,indent=2,ensure_ascii=False)+"\n")
def read_jsonl(path:Path)->list[dict[str,Any]]:
 with path.open(encoding="utf-8") as f:return [json.loads(x) for x in f if x.strip()]
def build_candidates(pool_size:int,target:int)->dict[str,Any]:
 import pyarrow.parquet as pq
 if pool_size<target:raise ValueError("pool_size must be >= target")
 shards=sorted((ROOT/"metadata").glob("train-*-of-00005.parquet"))
 if len(shards)!=5:raise RuntimeError(f"expected five shards, got {len(shards)}")
 heap=[];seq=0;totals=Counter()
 cols=["song_id","title","artist_names","artist_ids","album_name","album_id","isExplicit","views","duration"]
 for shard_no,shard in enumerate(shards):
  pf=pq.ParquetFile(shard);totals["metadata_rows"]+=pf.metadata.num_rows
  for rg in range(pf.metadata.num_row_groups):
   for raw in pf.read_row_group(rg,columns=cols).to_pylist():
    totals["rows_scanned"]+=1;sid=str(raw.get("song_id") or "").strip()
    try:duration=int(raw.get("duration"))
    except (TypeError,ValueError):totals["invalid_duration"]+=1;continue
    if duration<45 or duration>360:totals["duration_filtered"]+=1;continue
    if not ID_RE.fullmatch(sid):totals["invalid_song_id"]+=1;continue
    rank=stable_rank(sid)
    rec={"rank_sha256":rank,"song_id":sid,"url":"https://www.youtube.com/watch?v="+sid,"title":raw.get("title") or "",
     "artist_names":raw.get("artist_names") or [],"artist_ids":raw.get("artist_ids") or [],
     "album_name":raw.get("album_name") or "","album_id":raw.get("album_id") or "",
     "is_explicit":bool(raw.get("isExplicit")) if raw.get("isExplicit") is not None else None,
     "views":raw.get("views") or "","metadata_duration_seconds":duration,"source_shard":shard_no}
    item=(-int(rank,16),sid,seq,rec);seq+=1
    if len(heap)<pool_size:heapq.heappush(heap,item)
    elif item[0]>heap[0][0]:heapq.heapreplace(heap,item)
 ranked=sorted((x[3] for x in heap),key=lambda x:(x["rank_sha256"],x["song_id"]))
 accepted=[];seen_song=set();seen_album=set();artist_count=Counter()
 for rec in ranked:
  sid=rec["song_id"]
  if sid in seen_song:totals["duplicate_song_id"]+=1;continue
  seen_song.add(sid);ak=album_key(rec);artists=artist_keys(rec)
  if ak in seen_album:totals["album_cap_rejected"]+=1;continue
  if any(artist_count[a]>=2 for a in artists):totals["artist_cap_rejected"]+=1;continue
  rec["album_cap_key"]=ak;rec["artist_cap_keys"]=artists;rec["candidate_index"]=len(accepted);accepted.append(rec);seen_album.add(ak)
  for a in artists:artist_count[a]+=1
  if len(accepted)==target:break
 if len(accepted)!=target:raise RuntimeError(f"only {len(accepted)} candidates after caps; increase pool_size")
 atomic_text(CANDIDATES,"".join(json.dumps(x,sort_keys=True,ensure_ascii=False,separators=(",",":"))+"\n" for x in accepted))
 summary={"schema_version":1,"dataset":"laion/LAION-DISCO-12M","revision":REVISION,
  "selection":{"rank":"ascending sha256(revision + NUL + song_id)","metadata_duration_seconds_inclusive":[45,360],
   "album_cap":1,"artist_cap":2,"cap_scope":"all listed artist IDs; normalized names when IDs absent",
   "pool_size":pool_size,"candidate_count":target,"intended_success_count":2000},
  "counts":dict(totals),"candidate_manifest_sha256":sha256_file(CANDIDATES),
  "first_rank":accepted[0]["rank_sha256"],"last_rank":accepted[-1]["rank_sha256"]}
 write_json(CANDIDATE_SUMMARY,summary);return summary
def validate_candidates()->dict[str,Any]:
 rows=read_jsonl(CANDIDATES)
 if len(rows)<2000:raise AssertionError("candidate pool cannot support 2,000 successes even with zero failures")
 if len({r["song_id"] for r in rows})!=len(rows):raise AssertionError("duplicate song_id")
 ordered=[(r["rank_sha256"],r["song_id"]) for r in rows]
 if ordered!=sorted(ordered):raise AssertionError("not rank sorted")
 albums=Counter(r["album_cap_key"] for r in rows)
 if max(albums.values(),default=0)>1:raise AssertionError("album cap violated")
 artists=Counter(a for r in rows for a in r["artist_cap_keys"])
 if max(artists.values(),default=0)>2:raise AssertionError("artist cap violated")
 for i,r in enumerate(rows):
  if r["candidate_index"]!=i:raise AssertionError("candidate index mismatch")
  if stable_rank(r["song_id"])!=r["rank_sha256"]:raise AssertionError("rank mismatch")
  if not 45<=int(r["metadata_duration_seconds"])<=360:raise AssertionError("duration filter violated")
  if r["url"]!="https://www.youtube.com/watch?v="+r["song_id"]:raise AssertionError("URL mismatch")
 out={"status":"PASS","candidate_count":len(rows),"unique_albums":len(albums),"unique_artist_keys":len(artists),
  "max_artist_uses":max(artists.values(),default=0),"manifest_sha256":sha256_file(CANDIDATES),"bytes":CANDIDATES.stat().st_size}
 print(json.dumps(out,sort_keys=True));return out
def probe(path:Path)->dict[str,Any]:
 cp=subprocess.run([FFPROBE,"-v","error","-select_streams","a:0","-show_entries",
  "stream=codec_name,codec_long_name,sample_fmt,sample_rate,channels,channel_layout,bits_per_sample,duration:format=duration,size",
  "-of","json",str(path)],check=True,capture_output=True,text=True,timeout=60)
 return json.loads(cp.stdout)
def wav_stats(path:Path)->dict[str,Any]:
 pcm_hash=hashlib.sha256();n=ss=sall=peak=clipped=diff=ls=rs=0
 with wave.open(str(path),"rb") as w:
  channels,rate,width,fc,ctype=w.getnchannels(),w.getframerate(),w.getsampwidth(),w.getnframes(),w.getcomptype()
  if (channels,rate,width,ctype)!=(2,44100,2,"NONE"):raise ValueError(f"not canonical PCM16 stereo 44.1k: {(channels,rate,width,ctype)}")
  while data:=w.readframes(65536):
   pcm_hash.update(data);vals=array.array("h");vals.frombytes(data)
   if sys.byteorder!="little":vals.byteswap()
   n+=len(vals)
   for i in range(0,len(vals),2):
    l,r=vals[i],vals[i+1];ls+=l;rs+=r;diff+=(l-r)*(l-r)
   for v in vals:
    av=abs(v);peak=max(peak,av);clipped+=int(av>=32767);sall+=v;ss+=v*v
 if n==0:raise ValueError("empty canonical audio")
 fc=n//2
 return {"pcm_sha256":pcm_hash.hexdigest(),"sample_rate":44100,"channels":2,"sample_width_bytes":2,"frame_count":fc,
  "duration_seconds":fc/44100.0,"peak_linear":peak/32768.0,"rms_linear":math.sqrt(ss/n)/32768.0,
  "dc_linear":(sall/n)/32768.0,"left_dc_linear":(ls/fc)/32768.0,"right_dc_linear":(rs/fc)/32768.0,
  "stereo_difference_rms_linear":math.sqrt(diff/fc)/32768.0,"clipped_sample_fraction":clipped/n}
def canonicalize(source:Path,dest:Path)->tuple[dict[str,Any],dict[str,Any]]:
 dest.parent.mkdir(parents=True,exist_ok=True)
 cp=subprocess.run([FFMPEG,"-nostdin","-hide_banner","-loglevel","error","-i",str(source),"-map","0:a:0","-vn","-sn","-dn",
  "-ac","2","-ar","44100","-c:a","pcm_s16le","-fflags","+bitexact","-flags:a","+bitexact","-map_metadata","-1","-y",str(dest)],
  capture_output=True,text=True,timeout=600)
 if cp.returncode:raise RuntimeError("ffmpeg: "+cp.stderr[-2000:])
 return probe(dest),wav_stats(dest)
def classify_download_error(error:str)->str:
 text=(error or "").casefold()
 systemic=("http error 429","too many requests","rate limit","sign in to confirm","not a bot","authentication required")
 if any(x in text for x in systemic):return "SYSTEMIC_RATE_LIMIT_OR_AUTH"
 unavailable=("video unavailable","not available","private video","private","removed","copyright removal","geo","country","http error 403","members-only")
 if any(x in text for x in unavailable):return "SKIP_UNAVAILABLE"
 return "DOWNLOAD_ERROR"
def ytdlp_download(rec:dict[str,Any],attempts:int=3)->dict[str,Any]:
 idx,sid=rec["candidate_index"],rec["song_id"];work=TMP/f"{idx:05d}_{sid}"
 if work.exists():shutil.rmtree(work)
 work.mkdir(parents=True);output=work/"%(id)s.%(ext)s";error=""
 for attempt in range(attempts):
  cp=subprocess.run([str(YTDLP),"--no-playlist","--no-cache-dir","--no-mtime","--quiet","--no-warnings","--socket-timeout","30",
   "--retries","3","--fragment-retries","3","-f","bestaudio/best","-o",str(output),rec["url"]],
   capture_output=True,text=True,timeout=900)
  if cp.returncode==0:break
  error=(cp.stderr or cp.stdout)[-4000:]
  if attempt+1<attempts:time.sleep(2**attempt)
 else:
  shutil.rmtree(work,ignore_errors=True);category=classify_download_error(error);return {"ok":False,"stage":"skip_unavailable" if category=="SKIP_UNAVAILABLE" else "download","error_category":category,"error":error}
 files=[p for p in work.iterdir() if p.is_file() and not p.name.endswith((".part",".ytdl",".json"))]
 if len(files)!=1:
  names=[p.name for p in files];shutil.rmtree(work,ignore_errors=True);return {"ok":False,"stage":"download_output","error":f"expected 1 media file, got {names}"}
 raw=files[0];wavp=work/f"{sid}.canonical.wav"
 try:
  raw_probe=probe(raw);raw_sha=sha256_file(raw);canonical_probe,stats=canonicalize(raw,wavp);canonical_sha=sha256_file(wavp)
  if not 45.0<=stats["duration_seconds"]<=360.0:raise ValueError(f"actual duration {stats['duration_seconds']:.6f}s outside [45,360]")
  if stats["rms_linear"]<=0:raise ValueError("silent canonical audio")
  raw.unlink()
  return {"ok":True,"work":str(work),"temp_canonical":str(wavp),"raw_filename":raw.name,"raw_sha256":raw_sha,
   "canonical_sha256":canonical_sha,"raw_probe":raw_probe,"canonical_probe":canonical_probe,"audio_stats":stats}
 except Exception as exc:
  shutil.rmtree(work,ignore_errors=True);return {"ok":False,"stage":"canonicalize_validate","error":repr(exc)}
def load_known()->tuple[dict[str,str],dict[str,str]]:
 if not KNOWN_HASHES.exists():return {},{}
 obj=json.loads(KNOWN_HASHES.read_text(encoding="utf-8"));by_wav={};by_pcm={}
 for r in obj["files"]:by_wav[r["canonical_sha256"]]=r["source_label"];by_pcm[r["audio_stats"]["pcm_sha256"]]=r["source_label"]
 return by_wav,by_pcm
def append_results(rows:list[dict[str,Any]])->None:
 RESULTS.parent.mkdir(parents=True,exist_ok=True)
 with RESULTS.open("a",encoding="utf-8") as f:
  for r in rows:f.write(json.dumps(r,sort_keys=True,ensure_ascii=False,separators=(",",":"))+"\n")
  f.flush();os.fsync(f.fileno())
def systemic_failure(row:dict[str,Any])->bool:
 if row.get("accepted"):return False
 text=(str(row.get("error") or "")+" "+str(row.get("stage") or "")).casefold()
 terms=("http error 429","too many requests","rate limit","sign in to confirm","not a bot","authentication required")
 return any(x in text for x in terms)
def ingest(limit:int,workers:int,smoke:bool)->dict[str,Any]:
 if workers<1 or workers>4:raise ValueError("workers must be 1..4")
 candidates=read_jsonl(CANDIDATES);known_wav,known_pcm=load_known()
 history=read_jsonl(RESULTS) if RESULTS.exists() else [];attempted={x["song_id"] for x in history}
 accepted=[x for x in history if x.get("accepted")];seen_wav=dict(known_wav);seen_pcm=dict(known_pcm)
 for x in accepted:
  final=Path(x["canonical_path"])
  if not final.is_file():raise RuntimeError(f"accepted file missing: {final}")
  seen_wav[x["canonical_sha256"]]=x["song_id"];seen_pcm[x["audio_stats"]["pcm_sha256"]]=x["song_id"]
 successes=len(accepted);cursor=0;start=time.monotonic();base_successes=successes;base_attempts=len(history)
 checkpoints=[x for x in (250,500,1000,1500,2000) if x<=limit]
 while successes<limit:
  while cursor<len(candidates) and candidates[cursor]["song_id"] in attempted:cursor+=1
  if cursor>=len(candidates):raise RuntimeError(f"candidate list exhausted at {successes}/{limit}")
  next_checkpoint=next((x for x in checkpoints if x>successes),limit)
  remaining=min(limit-successes,next_checkpoint-successes);batch=[]
  while cursor<len(candidates) and len(batch)<min(workers,remaining):
   rec=candidates[cursor];cursor+=1
   if rec["song_id"] not in attempted:batch.append(rec);attempted.add(rec["song_id"])
  with ThreadPoolExecutor(max_workers=len(batch)) as ex:outcomes=list(ex.map(ytdlp_download,batch))
  rows=[]
  for rec,out in zip(batch,outcomes):
   base={"candidate_index":rec["candidate_index"],"rank_sha256":rec["rank_sha256"],"song_id":rec["song_id"],"url":rec["url"],
    "metadata_duration_seconds":rec["metadata_duration_seconds"],"attempted_at_unix":int(time.time()),
    "retention_policy":"canonical_only_after_validation","raw_retained":False}
   if not out["ok"]:rows.append({**base,**out,"accepted":False});continue
   dup=seen_wav.get(out["canonical_sha256"]) or seen_pcm.get(out["audio_stats"]["pcm_sha256"])
   temp=Path(out.pop("temp_canonical"));work=Path(out.pop("work"))
   if dup:
    shutil.rmtree(work,ignore_errors=True);rows.append({**base,**out,"accepted":False,"stage":"exact_dedup","duplicate_of":dup,"error":"exact canonical WAV or PCM duplicate"});continue
   dest=CANONICAL/f"{rec['candidate_index']:05d}_{rec['song_id']}.wav";CANONICAL.mkdir(parents=True,exist_ok=True)
   os.replace(temp,dest);shutil.rmtree(work,ignore_errors=True)
   seen_wav[out["canonical_sha256"]]=rec["song_id"];seen_pcm[out["audio_stats"]["pcm_sha256"]]=rec["song_id"];successes+=1
   rows.append({**base,**out,"accepted":True,"canonical_path":str(dest.resolve()),"canonical_bytes":dest.stat().st_size})
  append_results(rows);history.extend(rows)
  recent=history[-200:]
  if len(recent)==200 and sum(bool(x.get("accepted")) for x in recent)<100:
   stop={"status":"STOPPED","reason":"rolling_200_acceptance_below_50_percent","accepted_total":successes,
    "attempted_total":len(history),"rolling_200_accepted":sum(bool(x.get("accepted")) for x in recent)}
   write_json(ROOT/"manifests/tranche.stopped.json",stop);raise RuntimeError(json.dumps(stop,sort_keys=True))
  consecutive=0
  for x in reversed(history):
   if systemic_failure(x):consecutive+=1
   else:break
  tail=history[-20:];systemic_tail=sum(systemic_failure(x) for x in tail)
  if consecutive>=8 or (len(tail)==20 and systemic_tail>=16 and not any(x.get("accepted") for x in tail)):
   stop={"status":"STOPPED","reason":"systemic_rate_limit_or_auth_failure","accepted_total":successes,
    "attempted_total":len(history),"consecutive_systemic_failures":consecutive,"systemic_failures_last_20":systemic_tail}
   write_json(ROOT/"manifests/tranche.stopped.json",stop);raise RuntimeError(json.dumps(stop,sort_keys=True))
  if successes in checkpoints:
   elapsed=max(time.monotonic()-start,1e-9);new_successes=successes-base_successes;rate_s=new_successes/elapsed
   checkpoint={"event":"checkpoint","accepted_total":successes,"attempted_total":len(history),
    "new_accepted":new_successes,"new_attempts":len(history)-base_attempts,
    "acceptance_rate_total":successes/len(history),"accepted_per_minute":rate_s*60,
    "eta_seconds":(limit-successes)/rate_s if rate_s>0 else None,
    "failure_counts":dict(Counter(x.get("stage","unknown") for x in history if not x.get("accepted"))),
    "canonical_bytes":sum(p.stat().st_size for p in CANONICAL.glob("*.wav")),"elapsed_seconds":elapsed}
   with (ROOT/"logs/checkpoints.jsonl").open("a",encoding="utf-8") as f:f.write(json.dumps(checkpoint,sort_keys=True)+"\n")
   print(json.dumps(checkpoint,sort_keys=True),flush=True)
 allrows=read_jsonl(RESULTS)
 summary={"status":"PASS","mode":"smoke" if smoke else "tranche","target_successes":limit,
  "accepted_total":sum(bool(x.get("accepted")) for x in allrows),"attempted_total":len(allrows),
  "failure_counts":dict(Counter(x.get("stage","accepted") for x in allrows if not x.get("accepted"))),
  "results_sha256":sha256_file(RESULTS),"canonical_bytes":sum(p.stat().st_size for p in CANONICAL.glob("*.wav")),
  "known_hash_count":len(known_pcm),"workers":workers}
 write_json(ROOT/("manifests/smoke.summary.json" if smoke else "manifests/tranche.summary.json"),summary);return summary
def build_known(inputs_manifest:Path)->dict[str,Any]:
 obj=json.loads(inputs_manifest.read_text(encoding="utf-8"));rows=[];work=TMP/"known_hashes"
 if work.exists():shutil.rmtree(work)
 work.mkdir(parents=True)
 try:
  for i,item in enumerate(obj["files"]):
   source=Path(item["path"])
   if not source.is_file():raise FileNotFoundError(source)
   dest=work/f"{i:03d}.wav";canonical_probe,stats=canonicalize(source,dest)
   rows.append({"source_label":item["source_label"],"raw_sha256":sha256_file(source),"canonical_sha256":sha256_file(dest),
    "canonical_probe":canonical_probe,"audio_stats":stats})
 finally:shutil.rmtree(work,ignore_errors=True)
 if len({x["source_label"] for x in rows})!=len(rows):raise AssertionError("duplicate source label")
 out={"schema_version":1,"pipeline":"ffmpeg -ac 2 -ar 44100 -c:a pcm_s16le; metadata stripped; bitexact flags","file_count":len(rows),"files":rows}
 write_json(KNOWN_HASHES,out);out["manifest_sha256"]=sha256_file(KNOWN_HASHES);return out
def main()->None:
 ap=argparse.ArgumentParser();sub=ap.add_subparsers(dest="cmd",required=True)
 b=sub.add_parser("build-candidates");b.add_argument("--pool-size",type=int,default=250000);b.add_argument("--target",type=int,default=20000)
 sub.add_parser("validate-candidates");k=sub.add_parser("build-known-hashes");k.add_argument("inputs_manifest",type=Path)
 s=sub.add_parser("smoke");s.add_argument("--successes",type=int,default=3)
 r=sub.add_parser("run");r.add_argument("--successes",type=int,default=2000);r.add_argument("--workers",type=int,default=4)
 a=ap.parse_args()
 if a.cmd=="build-candidates":print(json.dumps(build_candidates(a.pool_size,a.target),indent=2,sort_keys=True))
 elif a.cmd=="validate-candidates":validate_candidates()
 elif a.cmd=="build-known-hashes":print(json.dumps(build_known(a.inputs_manifest),indent=2,sort_keys=True))
 elif a.cmd=="smoke":print(json.dumps(ingest(a.successes,min(a.successes,3),True),indent=2,sort_keys=True))
 elif a.cmd=="run":print(json.dumps(ingest(a.successes,a.workers,False),indent=2,sort_keys=True))
if __name__=="__main__":main()