# retoor <retoor@molodetz.nl>
import argparse
import hashlib
import json
import mimetypes
from pathlib import Path

RESULT_NAME = "result.md"
PROGRESS_NAME = "progress.jsonl"
PROMPT_NAME = "prompt.md"
MANIFEST_NAME = "manifest.json"
OPERATIONAL_FILES = frozenset({RESULT_NAME, PROGRESS_NAME, PROMPT_NAME, MANIFEST_NAME})
DEFAULT_MAX_FILES = 500
DEFAULT_MAX_MB = 256


def flat_name(root, path):
    return path.relative_to(root).as_posix().replace("/", "__")


def digest(path):
    found = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(65536), b""):
            found.update(chunk)
    return found.hexdigest()


def assemble(root, max_files=DEFAULT_MAX_FILES, max_bytes=DEFAULT_MAX_MB * 1024 * 1024):
    root = Path(root)
    candidates = sorted(
        (path for path in root.rglob("*") if path.is_file() and flat_name(root, path) not in OPERATIONAL_FILES),
        key=lambda path: flat_name(root, path),
    )
    files, omitted, seen = [], [], set()
    total = 0
    for path in candidates:
        name = flat_name(root, path)
        stem, tail = name, 2
        while stem in seen:
            stem = f"{name}~{tail}"
            tail += 1
        seen.add(stem)
        try:
            size = path.stat().st_size
        except OSError:
            continue
        if len(files) >= max_files or total + size > max_bytes:
            omitted.append(stem)
            continue
        try:
            sha = digest(path)
        except OSError:
            continue
        mime = mimetypes.guess_type(stem)[0] or "application/octet-stream"
        files.append({"name": stem, "bytes": size, "sha256": sha, "mime": mime})
        total += size
    note = ""
    if omitted:
        shown = ", ".join(omitted[:20])
        extra = f" and {len(omitted) - 20} more" if len(omitted) > 20 else ""
        note = f"Omitted {len(omitted)} file(s) beyond budget caps: {shown}{extra}"
    return {
        "result_present": (root / RESULT_NAME).is_file(),
        "files": files,
        "total_bytes": total,
        "file_count": len(files),
        "omitted": omitted,
        "note": note,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Assemble manifest.json for a result directory.")
    parser.add_argument("root", help="Result directory to inventory")
    parser.add_argument("--max-files", type=int, default=DEFAULT_MAX_FILES, help="Published files cap")
    parser.add_argument("--max-mb", type=int, default=DEFAULT_MAX_MB, help="Published megabytes cap")
    parser.add_argument("--write", action="store_true", help="Write manifest.json into the directory")
    args = parser.parse_args()
    payload = assemble(args.root, max_files=args.max_files, max_bytes=args.max_mb * 1024 * 1024)
    if args.write:
        with open(Path(args.root) / MANIFEST_NAME, "w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2)
    print(json.dumps(payload))