#!/usr/bin/env python3
from __future__ import annotations
import hashlib, json, os, shutil, stat, wave
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

ROOT=Path("/home/ubuntu/minimax-laion-corpus")
LABEL="interim_tranche_678_due_systemic_bot_auth"
FINAL=ROOT/"versions"/LABEL
STAGING=ROOT/"versions"/("."+LABEL+".staging")
RESULTS=ROOT/"manifests/ingest-results.jsonl"
CANDIDATES=ROOT/"manifests/candidates.jsonl"
KNOWN=ROOT/"manifests/known_user22_hashes.json"
STOPPED=ROOT/"manifests/tranche.stopped.json"
CANONICAL_ROOT=(ROOT/"canonical").resolve()
SEED="music3lab.laion.interim678.split.v1"
REVISION="6e7bf3758a77301e46a715af894fefd79bb1da53"

def canon(obj):
 return json.dumps(obj,sort_keys=True,ensure_ascii=False,separators=(",",":")).encode()

def sha_bytes(data):
 return hashlib.sha256(data).hexdigest()

def sha_file(path,chunk=8<<20):
 h=hashlib.sha256()
 with path.open("rb") as f:
  while b:=f.read(chunk):h.update(b)
 return h.hexdigest()

def pcm_info(path):
 h=hashlib.sha256()
 with wave.open(str(path),"rb") as w:
  info={"frame_count":w.getnframes(),"sample_rate":w.getframerate(),"channels":w.getnchannels(),
        "sample_width_bytes":w.getsampwidth(),"compression":w.getcomptype()}
  if (info["sample_rate"],info["channels"],info["sample_width_bytes"],info["compression"])!=(44100,2,2,"NONE"):
   raise RuntimeError(f"noncanonical WAV {path}: {info}")
  while data:=w.readframes(65536):h.update(data)
 info["pcm_sha256"]=h.hexdigest()
 return info

def copy_verify(item):
 row,record=item
 src=Path(row["canonical_path"]).resolve()
 if src.parent!=CANONICAL_ROOT:raise RuntimeError(f"source outside canonical root: {src}")
 if not src.is_file():raise FileNotFoundError(src)
 if sha_file(src)!=record["canonical_sha256"]:raise RuntimeError(f"source WAV hash mismatch: {src}")
 srcinfo=pcm_info(src)
 for key in ("frame_count","sample_rate","channels","pcm_sha256"):
  if srcinfo[key]!=record[key]:raise RuntimeError(f"source {key} mismatch: {src}")
 dst=STAGING/record["relative_path"]
 dst.parent.mkdir(parents=True,exist_ok=True)
 shutil.copyfile(src,dst)
 if dst.stat().st_size!=record["canonical_bytes"]:raise RuntimeError(f"copy size mismatch: {dst}")
 if sha_file(dst)!=record["canonical_sha256"]:raise RuntimeError(f"copy WAV hash mismatch: {dst}")
 dstinfo=pcm_info(dst)
 for key in ("frame_count","sample_rate","channels","pcm_sha256"):
  if dstinfo[key]!=record[key]:raise RuntimeError(f"copy {key} mismatch: {dst}")
 os.chmod(dst,0o444)
 return record

def main():
 if FINAL.exists():raise FileExistsError(FINAL)
 if STAGING.exists():raise FileExistsError(STAGING)
 rows=[json.loads(x) for x in RESULTS.read_text(encoding="utf-8").splitlines() if x.strip()]
 accepted=sorted((x for x in rows if x.get("accepted")),key=lambda x:x["candidate_index"])
 if len(accepted)!=678:raise RuntimeError(f"expected 678 accepted, got {len(accepted)}")
 if len(rows)!=994:raise RuntimeError(f"expected 994 attempts, got {len(rows)}")
 if [x["candidate_index"] for x in rows]!=list(range(994)):raise RuntimeError("attempt order is not frozen candidate order 0..993")
 if len({x["song_id"] for x in accepted})!=678:raise RuntimeError("duplicate source ID")
 wav_hashes=[x["canonical_sha256"] for x in accepted]
 pcm_hashes=[x["audio_stats"]["pcm_sha256"] for x in accepted]
 if len(set(wav_hashes))!=678 or len(set(pcm_hashes))!=678:raise RuntimeError("within-tranche exact duplicate")
 known=json.loads(KNOWN.read_text(encoding="utf-8"))["files"]
 if set(wav_hashes)&{x["canonical_sha256"] for x in known}:raise RuntimeError("user22 WAV overlap")
 if set(pcm_hashes)&{x["audio_stats"]["pcm_sha256"] for x in known}:raise RuntimeError("user22 PCM overlap")
 candidate_sha=sha_file(CANDIDATES);results_sha=sha_file(RESULTS);stopped_sha=sha_file(STOPPED)
 if candidate_sha!="533d8623268b18f87de2774ecf45656520f37edd8d5bb30c69f9bb6401231a07":raise RuntimeError("candidate manifest changed")
 stop=json.loads(STOPPED.read_text())
 if stop.get("accepted_total")!=678 or stop.get("reason")!="systemic_rate_limit_or_auth_failure":raise RuntimeError("stop artifact mismatch")
 records=[]
 for row in accepted:
  stats=row["audio_stats"];source_identity={
   "candidate_index":row["candidate_index"],"candidate_rank_sha256":row["rank_sha256"],
   "source_id":row["song_id"],"canonical_sha256":row["canonical_sha256"],
   "pcm_sha256":stats["pcm_sha256"]}
  split_rank=sha_bytes((SEED+"\0"+str(source_identity["candidate_index"])+"\0"+
   source_identity["candidate_rank_sha256"]+"\0"+source_identity["source_id"]+"\0"+
   source_identity["canonical_sha256"]+"\0"+source_identity["pcm_sha256"]).encode())
  records.append({
   "schema_version":"music3lab.laion.interim-record.v1",
   "candidate_index":row["candidate_index"],"candidate_rank_sha256":row["rank_sha256"],
   "source_dataset":"laion/LAION-DISCO-12M","source_dataset_revision":REVISION,
   "source_id":row["song_id"],"source_url":row["url"],
   "source_record_sha256":sha_bytes(canon(row)),
   "canonical_sha256":row["canonical_sha256"],"pcm_sha256":stats["pcm_sha256"],
   "canonical_bytes":row["canonical_bytes"],"frame_count":stats["frame_count"],
   "sample_rate":stats["sample_rate"],"channels":stats["channels"],
   "relative_path":"files/"+row["canonical_sha256"]+".wav","split_rank_sha256":split_rank,
  })
 by_split_rank=sorted(records,key=lambda x:(x["split_rank_sha256"],x["source_id"]))
 for i,record in enumerate(by_split_rank):
  record["split"]="train" if i<542 else ("validation" if i<610 else "heldout")
 counts={x:sum(r["split"]==x for r in records) for x in ("train","validation","heldout")}
 if counts!={"train":542,"validation":68,"heldout":68}:raise RuntimeError(counts)
 for split in counts:
  ids=[r["source_id"] for r in records if r["split"]==split]
  if len(ids)!=len(set(ids)):raise RuntimeError(f"duplicate within {split}")
 sets={s:{r["source_id"] for r in records if r["split"]==s} for s in counts}
 if sets["train"]&sets["validation"] or sets["train"]&sets["heldout"] or sets["validation"]&sets["heldout"]:
  raise RuntimeError("source leakage across splits")
 STAGING.mkdir(mode=0o755)
 (STAGING/"files").mkdir(mode=0o755)
 row_by_id={x["song_id"]:x for x in accepted}
 with ThreadPoolExecutor(max_workers=4) as ex:
  copied=list(ex.map(copy_verify,((row_by_id[r["source_id"]],r) for r in records)))
 if len(copied)!=678:raise RuntimeError("copy count mismatch")
 file_paths=list((STAGING/"files").glob("*.wav"))
 if len(file_paths)!=678:raise RuntimeError(f"copy file count {len(file_paths)}")
 manifest_rows=sorted(records,key=lambda x:x["candidate_index"])
 manifest_payload=b"".join(canon(x)+b"\n" for x in manifest_rows)
 semantic_payload={"schema_version":"music3lab.laion.interim-semantic.v1","label":LABEL,
  "complete_corpus":False,"incomplete_reason":"systemic_bot_auth","original_target":2000,
  "source_dataset":"laion/LAION-DISCO-12M","source_dataset_revision":REVISION,
  "candidate_manifest_sha256":candidate_sha,"stop_artifact_sha256":stopped_sha,
  "split_seed":SEED,"split_rank":"sha256(seed NUL candidate_index NUL candidate_rank NUL source_id NUL canonical_sha256 NUL pcm_sha256)",
  "split_assignment":"ascending split_rank: first 542 train, next 68 validation, final 68 heldout",
  "records":[{k:r[k] for k in ("candidate_index","candidate_rank_sha256","source_id","canonical_sha256","pcm_sha256",
    "canonical_bytes","frame_count","sample_rate","channels","relative_path","split_rank_sha256","split")} for r in manifest_rows]}
 semantic_digest=sha_bytes(canon(semantic_payload))
 files_payload="".join(f"{r['canonical_sha256']}  {r['canonical_bytes']}  {r['relative_path']}\n" for r in sorted(records,key=lambda x:x["relative_path"])).encode()
 files_digest=sha_bytes(files_payload)
 splits_obj={"schema_version":"music3lab.laion.interim-splits.v1","seed":SEED,
  "algorithm":semantic_payload["split_rank"],"assignment":semantic_payload["split_assignment"],
  "counts":counts,"splits":{s:[{"source_id":r["source_id"],"candidate_index":r["candidate_index"],
  "relative_path":r["relative_path"]} for r in sorted((x for x in records if x["split"]==s),key=lambda x:x["split_rank_sha256"])] for s in counts}}
 splits_payload=canon(splits_obj)+b"\n"
 manifest_path=STAGING/"manifest.jsonl";splits_path=STAGING/"splits.json";summary_path=STAGING/"summary.json"
 manifest_path.write_bytes(manifest_payload);splits_path.write_bytes(splits_payload)
 summary={"schema_version":"music3lab.laion.interim-summary.v1","label":LABEL,"status":"INTERIM_NOT_COMPLETE",
  "incomplete_reason":"systemic_bot_auth","original_target":2000,"accepted_count":678,"attempted_count":994,
  "skip_unavailable_count":299,"systemic_failure_count":17,"next_candidate_index":994,
  "source_dataset":"laion/LAION-DISCO-12M","source_dataset_revision":REVISION,
  "candidate_manifest_sha256":candidate_sha,"ingest_results_sha256":results_sha,"stop_artifact_sha256":stopped_sha,
  "manifest_sha256":sha_file(manifest_path),"splits_sha256":sha_file(splits_path),
  "semantic_digest_sha256":semantic_digest,"files_aggregate_digest_sha256":files_digest,
  "file_count":678,"canonical_bytes":sum(r["canonical_bytes"] for r in records),
  "split_seed":SEED,"split_counts":counts,"within_wav_duplicates":0,"within_pcm_duplicates":0,
  "user22_wav_overlap":0,"user22_pcm_overlap":0}
 summary_path.write_bytes(canon(summary)+b"\n")
 sums=[]
 for r in sorted(records,key=lambda x:x["relative_path"]):sums.append(f"{r['canonical_sha256']}  {r['relative_path']}\n")
 for name in ("manifest.jsonl","splits.json","summary.json"):sums.append(f"{sha_file(STAGING/name)}  {name}\n")
 sums_path=STAGING/"SHA256SUMS";sums_path.write_text("".join(sums),encoding="utf-8")
 for p in (manifest_path,splits_path,summary_path,sums_path):os.chmod(p,0o444)
 os.chmod(STAGING/"files",0o555)
 os.chmod(STAGING,0o555)
 os.replace(STAGING,FINAL)
 final_summary=json.loads((FINAL/"summary.json").read_text())
 if final_summary!=summary:raise RuntimeError("published summary mismatch")
 for r in manifest_rows:
  p=FINAL/r["relative_path"]
  if not p.is_file() or sha_file(p)!=r["canonical_sha256"]:raise RuntimeError(f"postpublish mismatch {p}")
  if stat.S_IMODE(p.stat().st_mode)!=0o444:raise RuntimeError(f"mode mismatch {p}")
 print(json.dumps({"status":"PASS","root":str(FINAL),"manifest":str(FINAL/"manifest.jsonl"),
  "manifest_sha256":summary["manifest_sha256"],"summary_sha256":sha_file(FINAL/"summary.json"),
  "splits_sha256":summary["splits_sha256"],"semantic_digest_sha256":semantic_digest,
  "files_aggregate_digest_sha256":files_digest,"sha256sums_sha256":sha_file(FINAL/"SHA256SUMS"),
  "accepted_count":678,"canonical_bytes":summary["canonical_bytes"],"split_counts":counts},sort_keys=True))
if __name__=="__main__":main()