"""d02_audio_video_pass.py (device side). Header and signal statistics for every WAV,
ffprobe metadata for every MP4, and a size/mtime/sha256 inventory of every file.

WAV: channels, sample width, rate, frames, duration; DC offset, std, extremes, number of
samples at full scale (clipping check); 10 Hz block-RMS envelope saved for synchronisation checks.
sha256 is computed for files under 50 MB only (labels, MTConnect, CT, NC, XML, summary);
large files are identified by size and mtime. Resumable with a time budget (argv[1], s).
"""
import sys, os, json, time, glob, wave, hashlib, subprocess
import numpy as np

ROOT = os.path.expanduser("~/mnt/archive (1)")
OUT = os.path.join(ROOT, "Phase2", "derived", "audio_stats")
budget = float(sys.argv[1]) if len(sys.argv) > 1 else 140.0
t_start = time.time()
os.makedirs(OUT, exist_ok=True)

def runkey(p):
    parts = p.split(os.sep)
    return (parts[-3], int(parts[-2].replace("dataset", "")), parts[-1])

wavs = sorted(glob.glob(os.path.join(ROOT, "*", "dataset*", "*.wav")), key=runkey)
for f in wavs:
    mach, run, name = f.split(os.sep)[-3], f.split(os.sep)[-2], os.path.basename(f)
    ch = name.rsplit("_", 1)[-1].replace(".wav", "")
    tag = f"{mach}_{run}_{ch}"
    js = os.path.join(OUT, tag + ".json")
    if os.path.exists(js):
        continue
    size = os.path.getsize(f)
    if time.time() - t_start + size / 60e6 + 3 > budget:
        break
    t0 = time.time()
    with wave.open(f, "rb") as w:
        nch, sw, fr, nfr, comp = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes(), w.getcomptype()
        raw = w.readframes(nfr)
    rec = {"machine": mach, "run": run, "channel": ch, "file": name, "bytes": size,
           "nchannels": nch, "sampwidth": sw, "framerate": fr, "nframes": nfr, "comptype": comp,
           "duration_s": nfr / fr}
    if sw == 2:
        x = np.frombuffer(raw, dtype="<i2")
        if nch > 1:
            x = x.reshape(-1, nch)[:, 0]
        xf = x.astype(np.float32)
        rec.update({"mean": float(xf.mean()), "std": float(xf.std()), "min": int(x.min()), "max": int(x.max()),
                    "n_fullscale_pos": int((x == 32767).sum()), "n_fullscale_neg": int((x == -32768).sum()),
                    "n_zero": int((x == 0).sum())})
        B = fr // 10
        m = (len(xf) // B) * B
        blk = xf[:m].reshape(-1, B)
        env = np.sqrt(((blk - blk.mean(axis=1, keepdims=True)) ** 2).mean(axis=1)).astype(np.float32)
        np.save(os.path.join(OUT, tag + "_env10hz.npy"), env)
        rec["env_rows"] = int(env.shape[0])
    rec["seconds_to_process"] = round(time.time() - t0, 1)
    with open(js, "w") as fh:
        json.dump(rec, fh, indent=1)
left = sum(1 for f in wavs if not os.path.exists(os.path.join(OUT, f"{f.split(os.sep)[-3]}_{f.split(os.sep)[-2]}_{os.path.basename(f).rsplit('_',1)[-1].replace('.wav','')}.json")))
print(f"wav remaining {left} of {len(wavs)}; elapsed {time.time()-t_start:.0f}s")

if left == 0:
    vid_js = os.path.join(ROOT, "Phase2", "derived", "video_probe.json")
    if not os.path.exists(vid_js):
        vids = []
        for f in sorted(glob.glob(os.path.join(ROOT, "*", "dataset*", "*.mp4")), key=runkey):
            r = subprocess.run(["ffprobe", "-v", "error", "-print_format", "json", "-show_format", "-show_streams", f],
                               capture_output=True, text=True)
            j = json.loads(r.stdout) if r.returncode == 0 else {}
            vs = [s for s in j.get("streams", []) if s.get("codec_type") == "video"]
            aus = [s for s in j.get("streams", []) if s.get("codec_type") == "audio"]
            v = vs[0] if vs else {}
            vids.append({"machine": f.split(os.sep)[-3], "run": f.split(os.sep)[-2], "file": os.path.basename(f),
                         "bytes": os.path.getsize(f), "ok": r.returncode == 0,
                         "duration_s": float(j.get("format", {}).get("duration", "nan")),
                         "codec": v.get("codec_name"), "width": v.get("width"), "height": v.get("height"),
                         "avg_frame_rate": v.get("avg_frame_rate"), "nb_frames": v.get("nb_frames"),
                         "n_audio_streams": len(aus), "creation_time": j.get("format", {}).get("tags", {}).get("creation_time")})
        with open(vid_js, "w") as fh:
            json.dump(vids, fh, indent=1)
        print("video probed", len(vids))
    inv_js = os.path.join(ROOT, "Phase2", "derived", "file_inventory.json")
    if not os.path.exists(inv_js) and time.time() - t_start < budget - 40:
        inv = []
        for dp, dn, fn in os.walk(ROOT):
            if os.sep + "Phase2" in dp or dp.endswith("Phase2"):
                continue
            for n in fn:
                p = os.path.join(dp, n)
                rel = os.path.relpath(p, ROOT)
                st = os.stat(p)
                rec = {"path": rel.replace(os.sep, "/"), "bytes": st.st_size, "mtime": st.st_mtime}
                if st.st_size < 50e6:
                    h = hashlib.sha256()
                    with open(p, "rb") as fh:
                        for chunk in iter(lambda: fh.read(1 << 20), b""):
                            h.update(chunk)
                    rec["sha256"] = h.hexdigest()
                inv.append(rec)
        inv.sort(key=lambda r: r["path"])
        with open(inv_js, "w") as fh:
            json.dump(inv, fh, indent=1)
        print("inventory files", len(inv), "bytes", sum(r["bytes"] for r in inv))