"""
Step 1: verified GF(256) multiply (RS keystone) + content-addressed chunk store.

  (a) LOG / EXP verified N/N (255/255 each)
  (b) COMPOSED GF(256) multiply verified over all 65,536 (a,b) pairs
  (c) chunk store reconstructs a 'drive image' bit-exact
  (d) dedup removes redundancy (real savings), and is bit-verified by SHA-256
  (e) honesty: incompressible data yields ~no savings (no beating entropy)
"""
import os
import torch
from storage import gf256
from storage.common import verify
from storage.chunkstore import ChunkStore

torch.manual_seed(0)

print("=" * 62)
print("STEP 1  --  verified GF(256) + content-addressed chunk store")
print("=" * 62)

# ---- verified coding keystone ----
net_log, net_exp = gf256.train_units()
lok, ltot = verify(net_log, *gf256.log_domain())
eok, etot = verify(net_exp, *gf256.exp_domain())
mok, mtot = gf256.verify_mul(net_log, net_exp)
print("-" * 62)
print(f"(a) LOG verified {lok}/{ltot}  |  EXP verified {eok}/{etot}  -> "
      f"{'PASS' if lok==ltot and eok==etot else 'FAIL'}")
print(f"(b) COMPOSED GF(256) multiply verified {mok}/{mtot}  -> "
      f"{'PASS' if mok==mtot else 'FAIL'}")

# ---- content-addressed chunk store on a synthetic 'drive image' ----
fileA = os.urandom(64 * 1024)
fileB = os.urandom(64 * 1024)
fileC = os.urandom(64 * 1024)
tail = os.urandom(64 * 1024)
# a drive with duplicated files (A and B appear twice) + unique C + unique tail
image = fileA + fileB + fileA + fileC + fileB + tail

cs = ChunkStore()
manifest = cs.put_stream(image)
rebuilt = cs.reconstruct(manifest)

exact = rebuilt == image
print("-" * 62)
print(f"(c) reconstruct drive image: {'bit-exact PASS' if exact else 'FAIL'} "
      f"({len(image)} bytes)")
print(f"(d) integrity (all chunks hash-verified): "
      f"{'PASS' if cs.integrity_ok() else 'FAIL'}")
st = cs.stats()
print(f"    logical {st['logical_bytes']//1024} KB  ->  stored {st['stored_bytes']//1024} KB   "
      f"dedup {st['dedup_ratio']:.2f}x   "
      f"({st['unique_chunks']}/{st['total_chunks']} chunks unique)")

# ---- honesty check: incompressible random data must NOT shrink ----
rnd = os.urandom(384 * 1024)
cs2 = ChunkStore()
cs2.reconstruct(cs2.put_stream(rnd))  # store it
st2 = cs2.stats()
print(f"(e) incompressible 384 KB: dedup {st2['dedup_ratio']:.2f}x  "
      f"-> {'PASS (no magic)' if st2['dedup_ratio'] < 1.05 else 'unexpected'}")

allpass = (lok == ltot and eok == etot and mok == mtot and exact
           and cs.integrity_ok() and st['dedup_ratio'] > 1.2 and st2['dedup_ratio'] < 1.05)
print("=" * 62)
print(f"OVERALL: {'ALL PASS' if allpass else 'some checks failed'}")

if lok == ltot and eok == etot and mok == mtot:
    torch.save({"log": net_log.state_dict(), "exp": net_exp.state_dict(),
                "meta": {"unit": "GF256_MUL(log/exp)", "poly": "0x11D",
                         "multiply_verified": f"{mok}/{mtot}"}}, "GF256.pt")
    print("saved verified units -> GF256.pt")