from __future__ import annotations

import hashlib
from pathlib import Path


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return f"sha256:{h.hexdigest()}"


def sha256_bytes(data: bytes) -> str:
    return f"sha256:{hashlib.sha256(data).hexdigest()}"


def tree_hash(file_checksums: dict[str, str]) -> str:
    """Canonical tree hash per docs/protocol.md."""
    lines: list[bytes] = []
    for rel in sorted(file_checksums.keys()):
        digest = file_checksums[rel]
        if digest.startswith("sha256:"):
            hexdigest = digest[len("sha256:") :]
        else:
            hexdigest = digest
        lines.append(f"{rel}\0{hexdigest}\n".encode("utf-8"))
    return sha256_bytes(b"".join(lines))


def checksum_dir(root: Path, files: list[Path] | None = None) -> dict[str, str]:
    """Hash files under root; keys are relative posix paths."""
    out: dict[str, str] = {}
    if files is None:
        paths = sorted(p for p in root.rglob("*") if p.is_file())
    else:
        paths = files
    for p in paths:
        rel = p.relative_to(root).as_posix()
        out[rel] = sha256_file(p)
    return out