"""
Self-healing .pt disc archive.

A disc image is chunked (fixed size), duplicate chunks deduped, and each unique
chunk Reed-Solomon split into n = k+m shards (verified GF(256) core). Every shard
carries its own SHA-256, so SILENT corruption (bit-rot flipping bytes) is
DETECTED, marked as an erasure, and REPAIRED from the surviving shards -- as long
as at least k of the n shards per chunk are still intact.

The whole thing serializes to a single .pt (torch.save):
  { meta, manifest:[chunk_hash...], chunks:{ hash: {shards, shard_hashes, L, orig} } }

Honesty: RS adds redundancy (archive is ~(k+m)/k x the deduped image). It repairs
up to m damaged shards per chunk; beyond that a chunk is flagged as lost, never
silently wrong. No entropy is beaten.
"""
from __future__ import annotations
import hashlib
import time
import torch
from . import rs

CHUNK = 4096


def _h(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def build(data: bytes, label="disc", k=4, m=2, chunk=CHUNK, bad_sectors=None) -> dict:
    chunks = {}
    manifest = []
    for off in range(0, len(data), chunk):
        ch = data[off:off + chunk]
        h = _h(ch)
        manifest.append(h)
        if h in chunks:
            continue
        shards, L, orig = rs.encode(ch, k, m)
        shards = [bytes(s) for s in shards]
        chunks[h] = {"shards": shards, "shard_hashes": [_h(s) for s in shards],
                     "L": L, "orig": orig, "k": k, "m": m}
    return {
        "meta": {"label": label, "size": len(data), "chunk": chunk, "k": k, "m": m,
                 "bad_sectors": list(bad_sectors or []), "created": time.time(),
                 "format": "neural-cd-preserve/1"},
        "manifest": manifest, "chunks": chunks,
    }


def save(arc: dict, path: str):
    torch.save(arc, path)


def load(path: str) -> dict:
    return torch.load(path, map_location="cpu", weights_only=False)


def _good_shards(rec: dict) -> dict:
    """{index: bytes} for shards whose hash still matches (uncorrupted)."""
    out = {}
    for i, s in enumerate(rec["shards"]):
        if _h(s) == rec["shard_hashes"][i]:
            out[i] = s
    return out


def verify(arc: dict) -> dict:
    """Report health without modifying. -> {chunks, corrupted_shards, repairable,
    lost} where lost chunks have < k intact shards (unrecoverable)."""
    corrupted = repairable = lost = 0
    for h, rec in arc["chunks"].items():
        good = _good_shards(rec)
        nbad = len(rec["shards"]) - len(good)
        corrupted += nbad
        if nbad and len(good) >= rec["k"]:
            repairable += 1
        if len(good) < rec["k"]:
            lost += 1
    return {"chunks": len(arc["chunks"]), "corrupted_shards": corrupted,
            "repairable_chunks": repairable, "lost_chunks": lost,
            "healthy": corrupted == 0 and lost == 0}


def heal(arc: dict) -> int:
    """Detect corrupted shards and regenerate them from survivors. Returns the
    number of shards repaired. Chunks with < k intact shards cannot be healed."""
    repaired = 0
    for h, rec in arc["chunks"].items():
        good = _good_shards(rec)
        if len(good) == len(rec["shards"]):
            continue
        if len(good) < rec["k"]:
            continue                                    # unrecoverable, leave as-is
        ch = rs.decode(good, rec["k"], rec["m"], rec["L"], rec["orig"])
        if _h(ch) != h:
            continue                                    # decoded wrong -> don't trust
        fresh, _, _ = rs.encode(ch, rec["k"], rec["m"])
        for i in range(len(rec["shards"])):
            if i not in good:
                rec["shards"][i] = bytes(fresh[i])
                rec["shard_hashes"][i] = _h(rec["shards"][i])
                repaired += 1
    return repaired


def restore(arc: dict) -> bytes:
    """Reconstruct the disc image, RS-correcting any detected corruption."""
    out = bytearray()
    for h in arc["manifest"]:
        rec = arc["chunks"][h]
        good = _good_shards(rec)
        if len(good) < rec["k"]:
            raise IOError(f"chunk {h[:8]}: only {len(good)} intact shards, "
                          f"need {rec['k']} -- unrecoverable")
        ch = rs.decode(good, rec["k"], rec["m"], rec["L"], rec["orig"])
        if _h(ch) != h:
            raise IOError(f"chunk {h[:8]}: hash mismatch after decode")
        out += ch
    return bytes(out[:arc["meta"]["size"]])