CIS490/receiver/store.py

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import AsyncIterator


# ULID/UUID/host-id alphabet plus separators — strict on purpose.
_ID_RE = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")


def is_valid_id(s: str) -> bool:
    return bool(_ID_RE.match(s))


@dataclass(frozen=True)
class StoreResult:
    status: str               # "stored" | "already-present" | "conflict" | "sha-mismatch"
    sha256: str | None
    size_bytes: int | None
    existing_sha256: str | None = None


class EpisodeStore:
    def __init__(
        self,
        store_root: Path,
        incoming_root: Path,
        index_path: Path,
    ) -> None:
        self.store_root = store_root
        self.incoming_root = incoming_root
        self.index_path = index_path
        self.store_root.mkdir(parents=True, exist_ok=True)
        self.incoming_root.mkdir(parents=True, exist_ok=True)
        self.index_path.parent.mkdir(parents=True, exist_ok=True)
        self.index_path.touch(exist_ok=True)

    def final_path(self, host_id: str, episode_id: str) -> Path:
        return self.store_root / host_id / f"{episode_id}.tar.zst"

    def _hash_file(self, 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 h.hexdigest()

    async def ingest_stream(
        self,
        host_id: str,
        episode_id: str,
        expected_sha256: str,
        schema_version: int,
        body: AsyncIterator[bytes],
        max_bytes: int,
        commit: str | None = None,
    ) -> StoreResult:
        final = self.final_path(host_id, episode_id)
        if final.exists():
            existing = self._hash_file(final)
            # Drain the body so the client doesn't see a write error.
            async for _ in body:
                pass
            if existing == expected_sha256:
                return StoreResult(
                    status="already-present",
                    sha256=existing,
                    size_bytes=final.stat().st_size,
                )
            return StoreResult(
                status="conflict",
                sha256=None,
                size_bytes=None,
                existing_sha256=existing,
            )

        incoming_dir = self.incoming_root / host_id
        incoming_dir.mkdir(parents=True, exist_ok=True)
        partial = incoming_dir / f"{episode_id}.tar.zst.partial"

        h = hashlib.sha256()
        bytes_written = 0
        try:
            with partial.open("wb") as out:
                async for chunk in body:
                    bytes_written += len(chunk)
                    if bytes_written > max_bytes:
                        partial.unlink(missing_ok=True)
                        return StoreResult(
                            status="too-large",
                            sha256=None,
                            size_bytes=bytes_written,
                        )
                    h.update(chunk)
                    out.write(chunk)
            actual = h.hexdigest()
            if actual != expected_sha256:
                partial.unlink(missing_ok=True)
                return StoreResult(
                    status="sha-mismatch",
                    sha256=actual,
                    size_bytes=bytes_written,
                )

            final.parent.mkdir(parents=True, exist_ok=True)
            os.replace(partial, final)
            row = {
                "received_at_wall": datetime.now(timezone.utc).isoformat(),
                "host_id": host_id,
                "episode_id": episode_id,
                "sha256": actual,
                "size_bytes": bytes_written,
                "schema_version": schema_version,
            }
            if commit:
                row["commit"] = commit
            self._append_index(row)
            return StoreResult(status="stored", sha256=actual, size_bytes=bytes_written)
        except BaseException:
            partial.unlink(missing_ok=True)
            raise

    def _append_index(self, row: dict) -> None:
        line = json.dumps(row, sort_keys=True) + "\n"
        # Single write of one line is atomic on POSIX for sizes < PIPE_BUF.
        # Index rows are well under that.
        with self.index_path.open("a") as f:
            f.write(line)