storage/vault.py

"""
Step 3 backing store: the Vault.

Combines everything so far into an on-disk repository:
  * files are content-defined-chunked (Step 1 CDC)
  * unique chunks are stored once, addressed by SHA-256 (dedup, bit-verified)
  * each unique chunk is Reed-Solomon split into n = k+m shard files (Step 2),
    so any k shards rebuild it -- up to m shard files can be lost/corrupted

On-disk layout:
  <vault>/chunks/<sha256>/shard_0 .. shard_{n-1}, meta.json
  <vault>/files/<relpath>.json   (manifest: name, size, [chunk hashes])

Reads reconstruct on demand and RS-heal transparently; `heal()` regenerates any
missing shard files from the survivors. Nothing here beats entropy -- dedup
removes redundancy, RS adds redundancy for resilience.
"""
from __future__ import annotations
import os
import json
import hashlib
from .chunkstore import chunk_stream
from . import rs


class Vault:
    def __init__(self, path: str, k: int = 4, m: int = 2):
        self.path = path
        self.k, self.m = k, m
        self.cdir = os.path.join(path, "chunks")
        self.fdir = os.path.join(path, "files")
        os.makedirs(self.cdir, exist_ok=True)
        os.makedirs(self.fdir, exist_ok=True)

    # ---- chunks ----
    def _cd(self, h: str) -> str:
        return os.path.join(self.cdir, h)

    def _store_chunk(self, ch: bytes) -> str:
        h = hashlib.sha256(ch).hexdigest()
        d = self._cd(h)
        if os.path.exists(os.path.join(d, "meta.json")):
            return h                                    # dedup: already stored
        os.makedirs(d, exist_ok=True)
        shards, L, orig = rs.encode(ch, self.k, self.m)
        for i, s in enumerate(shards):
            with open(os.path.join(d, f"shard_{i}"), "wb") as f:
                f.write(bytes(s))
        json.dump({"L": L, "orig": orig, "k": self.k, "m": self.m},
                  open(os.path.join(d, "meta.json"), "w"))
        return h

    def _load_chunk(self, h: str) -> bytes:
        d = self._cd(h)
        meta = json.load(open(os.path.join(d, "meta.json")))
        n = meta["k"] + meta["m"]
        present = {}
        for i in range(n):
            p = os.path.join(d, f"shard_{i}")
            if os.path.exists(p):
                b = open(p, "rb").read()
                if len(b) == meta["L"]:
                    present[i] = b
        if len(present) < meta["k"]:
            raise IOError(f"chunk {h[:8]}: {len(present)} shards, need {meta['k']}")
        ch = rs.decode(present, meta["k"], meta["m"], meta["L"], meta["orig"])
        if hashlib.sha256(ch).hexdigest() != h:
            raise IOError(f"chunk {h[:8]}: hash mismatch after decode")
        return ch

    # ---- files ----
    def store_file(self, relpath: str, data: bytes) -> None:
        chunks = [self._store_chunk(data[o:o + l]) for o, l in chunk_stream(data)]
        mfp = os.path.join(self.fdir, relpath + ".json")
        os.makedirs(os.path.dirname(mfp) or ".", exist_ok=True)
        json.dump({"name": relpath, "size": len(data), "chunks": chunks}, open(mfp, "w"))

    def store_folder(self, src: str) -> None:
        for root, _, names in os.walk(src):
            for nm in names:
                fp = os.path.join(root, nm)
                rel = os.path.relpath(fp, src).replace("\\", "/")
                with open(fp, "rb") as f:
                    self.store_file(rel, f.read())

    def list_files(self):
        out = []
        for root, _, names in os.walk(self.fdir):
            for nm in names:
                if nm.endswith(".json"):
                    mf = json.load(open(os.path.join(root, nm)))
                    out.append((mf["name"], mf["size"]))
        return sorted(out)

    def read_file(self, relpath: str) -> bytes:
        mf = json.load(open(os.path.join(self.fdir, relpath + ".json")))
        data = b"".join(self._load_chunk(h) for h in mf["chunks"])
        return data[:mf["size"]]

    def export(self, dst: str) -> int:
        n = 0
        for rel, _ in self.list_files():
            data = self.read_file(rel)
            op = os.path.join(dst, rel)
            os.makedirs(os.path.dirname(op) or ".", exist_ok=True)
            with open(op, "wb") as f:
                f.write(data)
            n += 1
        return n

    # ---- resilience ----
    def heal(self) -> int:
        healed = 0
        for h in os.listdir(self.cdir):
            d = self._cd(h)
            mp = os.path.join(d, "meta.json")
            if not os.path.isfile(mp):
                continue
            meta = json.load(open(mp))
            n = meta["k"] + meta["m"]
            missing = [i for i in range(n) if not os.path.exists(os.path.join(d, f"shard_{i}"))]
            if missing:
                ch = self._load_chunk(h)                # decode from survivors
                shards, _, _ = rs.encode(ch, meta["k"], meta["m"])
                for i in missing:
                    with open(os.path.join(d, f"shard_{i}"), "wb") as f:
                        f.write(bytes(shards[i]))
                healed += len(missing)
        return healed

    def integrity_ok(self) -> bool:
        try:
            for rel, _ in self.list_files():
                self.read_file(rel)
            return True
        except IOError:
            return False

    def stats(self) -> dict:
        logical = sum(sz for _, sz in self.list_files())
        stored = 0
        chunks = 0
        for h in os.listdir(self.cdir):
            d = self._cd(h)
            if not os.path.isfile(os.path.join(d, "meta.json")):
                continue
            chunks += 1
            for f in os.listdir(d):
                if f.startswith("shard_"):
                    stored += os.path.getsize(os.path.join(d, f))
        return {"logical": logical, "stored": stored, "unique_chunks": chunks,
                "overhead": stored / max(1, logical)}