#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""Measure ddrescue on disposable regular files, never devices or user inputs."""

import argparse
import hashlib
import json
import os
import platform
import re
import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path

SECTOR = 512
SOURCE = b"A" * SECTOR + bytes(SECTOR) + b"C" * SECTOR + b"D" * SECTOR
SOURCE += b"E" * SECTOR + b"F" * SECTOR + b"G" * SECTOR + b"H" * SECTOR
STATUS = "?*/-+"
ARCHIVE_SHA256 = "2264622d309d6c87a1cfc19148292b8859a688e9bc02d4702f5cd4f288745542"


def require(condition, message):
    if not condition:
        raise ValueError(message)


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


def c_integer(token):
    require(bool(re.fullmatch(r"(?:0[xX][0-9a-fA-F]+|[0-9]+)", token)), "invalid map integer")
    return int(token, 16 if token.lower().startswith("0x") else 8 if len(token) > 1 and token[0] == "0" else 10)


def parse_map(text, extent=len(SOURCE)):
    """Parse complete maps produced by this experiment, not loose domain maps."""
    rows = [line.split() for line in text.splitlines() if line.strip() and not line.lstrip().startswith("#")]
    require(bool(rows), "empty map")
    require(len(rows[0]) == 3 and len(rows[0][1]) == 1 and rows[0][1] in "?*/-+FG", "invalid map status line")
    c_integer(rows[0][0])
    require(rows[0][2].isdigit(), "invalid current pass")
    result, end = [], 0
    for row in rows[1:]:
        require(len(row) == 3 and row[2] in STATUS and len(row[2]) == 1, "invalid map row")
        pos, size = c_integer(row[0]), c_integer(row[1])
        require(pos == end and size > 0 and pos + size <= extent, "gap, overlap or invalid extent")
        require(pos % SECTOR == 0 and size % SECTOR == 0, "unaligned fixture map")
        result.append({"offset": pos, "length": size, "status": row[2]})
        end = pos + size
    require(end == extent, "incomplete fixture map")
    return result


def map_sectors(blocks):
    return [block["status"] for block in blocks for _ in range(block["length"] // SECTOR)]


def evidence(image, map_text):
    require(len(image) == len(SOURCE), "unexpected output length")
    blocks = parse_map(map_text)
    matches = [image[i:i + SECTOR] == SOURCE[i:i + SECTOR] for i in range(0, len(SOURCE), SECTOR)]
    return {
        "imageBytes": len(image), "imageSha256": sha(image),
        "sectorSha256": [sha(image[i:i + SECTOR]) for i in range(0, len(image), SECTOR)],
        "matchesSourceBySector": matches, "matchingSourceBytes": sum(matches) * SECTOR,
        "zeroSectors": [i for i in range(len(image) // SECTOR) if image[i * SECTOR:(i + 1) * SECTOR] == bytes(SECTOR)],
        "mapText": map_text, "mapBlocks": blocks,
        "mapStatusBySector": map_sectors(blocks),
        "finishedBytes": sum(b["length"] for b in blocks if b["status"] == "+"),
    }


def run(command, cwd):
    proc = subprocess.run(command, cwd=cwd, capture_output=True, text=True, timeout=20,
                          env={**os.environ, "LC_ALL": "C", "TZ": "UTC"})
    return {"argv": [Path(command[0]).name, *command[1:]], "returncode": proc.returncode,
            "stdout": proc.stdout, "stderr": proc.stderr}


def expected_map(case_id):
    states = ["+"] * 8
    if case_id in {"fresh-bad-data", "reused-bad-data", "sparse-reused-bad-data"}:
        states[2] = "-"
    elif case_id == "fresh-bad-zero":
        states[1] = "-"
    elif case_id == "generated-fresh-bad-data":
        states[1] = states[2] = "?"
    elif case_id == "generated-reused-bad-data":
        states[1] = "?"
    return states


def expected_bytes(case_id):
    data = bytearray(SOURCE)
    if case_id in {"fresh-bad-data", "generated-fresh-bad-data"}:
        data[1024:1536] = bytes(SECTOR)
    if case_id in {"reused-bad-data", "sparse-reused-bad-data", "generated-reused-bad-data", "generated-sparse-reused-bad-data"}:
        data[1024:1536] = b"Z" * SECTOR
    if case_id in {"sparse-reused-healthy", "sparse-reused-bad-data", "resume-sparse-reused", "generated-sparse-reused-bad-data"}:
        data[512:1024] = b"Z" * SECTOR
    if case_id == "modified-after-completion":
        data[1536:2048] = b"X" * SECTOR
    return bytes(data)


IDS = {"fresh-bad-data", "reused-bad-data", "fresh-bad-zero", "fresh-healthy", "sparse-fresh-healthy",
       "sparse-reused-healthy", "sparse-reused-bad-data", "resume-sparse-reused", "modified-after-completion",
       "generated-fresh-bad-data", "generated-reused-bad-data", "generated-sparse-reused-bad-data"}


def stable_observations(receipt):
    """Exclude timestamps and raw map comments, not actual byte or coverage evidence."""
    keys = ("id", "imageBytes", "imageSha256", "sectorSha256", "matchesSourceBySector", "matchingSourceBytes",
            "zeroSectors", "mapBlocks", "mapStatusBySector", "finishedBytes")
    return [{**{k: c[k] for k in keys}, "operationStatus": c["operation"]["returncode"],
             "doneStatus": c["doneStatus"]["returncode"], "firstSectorDoneStatus": c["firstSectorDoneStatus"]["returncode"]}
            for c in receipt["cases"]]


def validate(receipt):
    cases = receipt["cases"]
    require(len(cases) == len(IDS) and {c["id"] for c in cases} == IDS, "missing or duplicate cases")
    require(receipt["source"]["sha256"] == sha(SOURCE), "wrong source")
    require(receipt["source"]["bytes"] == len(SOURCE), "wrong source extent")
    for case in cases:
        ident = case["id"]
        truth = expected_bytes(ident)
        require(case["imageSha256"] == sha(truth), ident + ": unexpected image")
        require(case["sectorSha256"] == [sha(truth[i:i + SECTOR]) for i in range(0, len(truth), SECTOR)], ident + ": sector hash")
        truth_matches = [truth[i:i + SECTOR] == SOURCE[i:i + SECTOR] for i in range(0, len(truth), SECTOR)]
        require(case["matchesSourceBySector"] == truth_matches, ident + ": source matches")
        require(case["matchingSourceBytes"] == sum(truth_matches) * SECTOR, ident + ": byte count")
        require(case["imageBytes"] == len(SOURCE), ident + ": image size")
        require(case["zeroSectors"] == [i for i in range(8) if truth[i * SECTOR:(i + 1) * SECTOR] == bytes(SECTOR)], ident + ": zeros")
        blocks = parse_map(case["mapText"])
        require(case["mapBlocks"] == blocks, ident + ": parsed blocks")
        states = map_sectors(blocks)
        require(states == expected_map(ident) == case["mapStatusBySector"], ident + ": map states")
        require(case["finishedBytes"] == states.count("+") * SECTOR, ident + ": finished count")
        require(case["operation"]["returncode"] == 0, ident + ": operation failed")
        require(case["doneStatus"]["returncode"] == int(any(s != "+" for s in states)), ident + ": done status")
        require(case["firstSectorDoneStatus"]["returncode"] == 0, ident + ": limited domain status")
    return True


def observe(ddrescue, ddrescuelog):
    versions = [run([str(binary), "--version"], None) for binary in (ddrescue, ddrescuelog)]
    require(all(v["returncode"] == 0 and v["stdout"].splitlines()[0].endswith("1.30") for v in versions), "requires GNU ddrescue and ddrescuelog 1.30")
    cases, copies = [], {}
    with tempfile.TemporaryDirectory(prefix="hesela-ddrescue-") as tmp:
        root = Path(tmp)
        (root / "source.img").write_bytes(SOURCE)
        for name, bad in (("bad-data", 2), ("bad-zero", 1)):
            lines = ["0 ? 1"] + [f"{i * SECTOR} {SECTOR} {'-' if i == bad else '+'}" for i in range(8)]
            (root / (name + ".map")).write_text("\n".join(lines) + "\n")

        def collect(ident, operation, output, mapfile, details):
            map_text = (root / mapfile).read_text()
            image = (root / output).read_bytes()
            copies[ident] = (image, map_text)
            item = {"id": ident, **details, "operation": operation,
                    "doneStatus": run([str(ddrescuelog), "--done-status", mapfile], root),
                    "firstSectorDoneStatus": run([str(ddrescuelog), "--done-status", "--size=512", mapfile], root),
                    **evidence(image, map_text)}
            cases.append(item)

        specifications = [
            ("fresh-bad-data", False, False, "bad-data"),
            ("reused-bad-data", True, False, "bad-data"),
            ("fresh-bad-zero", False, False, "bad-zero"),
            ("fresh-healthy", False, False, None),
            ("sparse-fresh-healthy", False, True, None),
            ("sparse-reused-healthy", True, True, None),
            ("sparse-reused-bad-data", True, True, "bad-data"),
        ]
        for ident, reused, sparse, fault in specifications:
            output, mapfile = ident + ".img", ident + ".map"
            if reused:
                (root / output).write_bytes(b"Z" * len(SOURCE))
            command = [str(ddrescue), "--quiet", "--sector-size=512", "--cluster-size=1", "--retry-passes=0"]
            if sparse:
                command.append("--sparse")
            if fault:
                command.append("--test-mode=" + fault + ".map")
            operation = run(command + ["source.img", output, mapfile], root)
            collect(ident, operation, output, mapfile, {"mode": "rescue", "destinationInitially": "Z-filled" if reused else "absent", "sparse": sparse, "simulatedBadSector": 2 if fault == "bad-data" else 1 if fault else None})

        # Reuse precisely the same image/map; sparse is now disabled and the source is fully readable.
        output, mapfile = "sparse-reused-bad-data.img", "sparse-reused-bad-data.map"
        operation = run([str(ddrescue), "--quiet", "--sector-size=512", "--cluster-size=1", "--retry-passes=1", "source.img", output, mapfile], root)
        collect("resume-sparse-reused", operation, output, mapfile, {"mode": "resume", "parentCase": "sparse-reused-bad-data", "sparse": False, "simulatedBadSector": None})

        for parent in ("fresh-bad-data", "reused-bad-data", "sparse-reused-bad-data"):
            image, _ = copies[parent]
            output, mapfile = "generated-" + parent + ".img", "generated-" + parent + ".map"
            (root / output).write_bytes(image)
            operation = run([str(ddrescue), "--quiet", "--sector-size=512", "--generate-mode", "source.img", output, mapfile], root)
            collect("generated-" + parent, operation, output, mapfile, {"mode": "generate", "parentCase": parent, "generatedMapIsApproximate": True})

        output, mapfile = "modified.img", "modified.map"
        image, map_text = copies["fresh-healthy"]
        modified = bytearray(image)
        modified[1536:2048] = b"X" * SECTOR
        (root / output).write_bytes(modified)
        (root / mapfile).write_text(map_text)
        collect("modified-after-completion", {"argv": ["fixture-sector-overwrite", "sector=3", "byte=0x58"], "returncode": 0, "stdout": "", "stderr": ""}, output, mapfile, {"mode": "post-copy-modification", "parentCase": "fresh-healthy"})
        require((root / "source.img").read_bytes() == SOURCE, "source unexpectedly modified")

    receipt = {
        "schemaVersion": 1, "experiment": "ddrescue-output-provenance", "generatedAt": datetime.now(timezone.utc).isoformat(),
        "evidenceType": "constructed-software-experiment", "faultMethod": "GNU ddrescue --test-mode; no physical device failures",
        "environment": {"python": platform.python_version(), "kernel": platform.release(), "system": platform.system(),
                        "machine": platform.machine(), "versions": versions,
                        "executableSha256": {p.name: sha(p.read_bytes()) for p in (ddrescue, ddrescuelog)}},
        "upstream": {"version": "1.30", "archiveUrl": "https://ftp.gnu.org/gnu/ddrescue/ddrescue-1.30.tar.lz", "archiveSha256": ARCHIVE_SHA256},
        "harnessSha256": sha(Path(__file__).read_bytes()),
        "source": {"bytes": len(SOURCE), "sectorBytes": SECTOR, "sha256": sha(SOURCE), "sectors": ["A", "00", "C", "D", "E", "F", "G", "H"], "encoding": "512 repetitions of each ASCII byte; 00 is 512 zero bytes"},
        "cases": cases,
        "limitations": ["Regular files only; software-simulated read errors.", "One upstream version and build; no hardware recovery rates, power-loss or timing claims.", "Map '+' describes rescue bookkeeping, not a digest comparison with the source.", "Reused sparse output intentionally violates the documented empty-or-zeroed destination precondition."]}
    validate(receipt)
    return receipt


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ddrescue", required=True, type=Path)
    parser.add_argument("--ddrescuelog", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    require(not args.output.exists(), "refusing to replace existing output")
    receipt = observe(args.ddrescue.resolve(), args.ddrescuelog.resolve())
    with args.output.open("x") as stream:
        json.dump(receipt, stream, indent=2)
        stream.write("\n")
    print(f"Validated {len(receipt['cases'])} cases; saved {args.output}")


if __name__ == "__main__":
    main()