Layer 3: Repository

class ArchiveRepository(DatasetHandle):
    """
    Repository for file archive operations.

    Content-addressed blob storage (raw bytes): a blob is stored once per
    checksum.

    Optionally, extracted text (by different origins) can be stored and
    retrieved. As well, other programs can write arbitrary additional data to
    the archive (such as pdf page thumbnails or their own file metadata) via
    [`put_data`][ArchiveRepository.put_data].

    Example:
        ```python
        archive = ArchiveRepository(dataset="my_data", uri="s3://bucket/dataset")

        # Archive a file
        checksum = archive.store("path/to/file.pdf")

        # Stream file contents by checksum
        for chunk in archive.stream(checksum):
            process(chunk)
        ```
    """

    def __init__(self, dataset: str, uri: Uri) -> None:
        super().__init__(dataset, uri)
        self._txts = get_store(
            self._store_uri, serialization_mode="auto", raise_on_nonexist=False
        )

    def exists(self, checksum: str) -> bool:
        """Check if blob exists for the given checksum."""
        return self._store.exists(path.ArchiveKey(checksum).blob)

    def iterate_checksums(self) -> Generator[str, None, None]:
        """Iterate the checksums of all blobs in the archive."""
        # exactly the checksum depth – a `put_data` path may end in "blob" too
        glob = f"*/*/*/*/{path.ARCHIVE_BLOB}"
        for key in self._store.iterate_keys(prefix=path.ARCHIVE, glob=glob):
            yield key.split("/")[-2]

    def stream(self, checksum: str) -> BytesGenerator:
        """Stream blob contents as bytes."""
        yield from self._store.stream(path.ArchiveKey(checksum).blob)

    def open(self, checksum: str) -> ContextManager[IO[bytes]]:
        """Get a file-like handle for reading."""
        return self._store.open(path.ArchiveKey(checksum).blob, mode=DEFAULT_MODE)

    def to_uri(self, checksum: str) -> str:
        return self._store.to_uri(path.ArchiveKey(checksum).blob)

    def local_path(self, checksum: str) -> ContextManager[Path]:
        """
        Get the local path to the blob.

        If storage is local, returns actual path. Otherwise, creates
        a temporary local copy that is cleaned up after context exit.
        """
        return self._store.local_path(path.ArchiveKey(checksum).blob)

    def store(
        self, uri: Uri, checksum: str | None = None, tag_updated: bool = True
    ) -> str:
        """
        Archive a file from a local or remote URI.

        Args:
            uri: Local or remote URI to the file
            checksum: Content hash (skip computation if provided)
            tag_updated: Touch the `archive/last_updated` tag

        Returns:
            checksum
        """
        checksum = self.store_blob(uri, checksum)
        if tag_updated:
            # Notify archive was updated
            self.touch()
        self.log.info(f"Archived `{mask_uri(uri)}` ({checksum})", checksum=checksum)
        return checksum

    def store_blob(self, uri: Uri, checksum: str | None = None) -> str:
        """
        Store bytes blob from given uri if it doesn't exist yet.

        Args:
            uri: Local or remote URI to the file
            checksum: Content hash (skip computation if provided)

        Returns:
            checksum
        """
        if checksum:
            validate_checksum(checksum)
        if checksum and self.exists(checksum):
            self.log.debug("Blob already exists, skipping", checksum=checksum)
            return checksum

        with open_virtual(uri, algorithm=CHECKSUM_ALGORITHM) as fh:
            if self.exists(fh.checksum):
                self.log.debug("Blob already exists, skipping", checksum=fh.checksum)
                return fh.checksum

            self.log.info(f"Storing blob `{fh.checksum}` ...", checksum=fh.checksum)
            self.write_blob(fh, fh.checksum, check_exists=False)
            return fh.checksum

    def write_blob(
        self, fh: BinaryIO, checksum: str | None = None, check_exists: bool = True
    ) -> str:
        """Write a blob from the given open file-handler"""
        if check_exists and checksum and self.exists(checksum):
            self.log.debug("Blob already exists, skipping", checksum=checksum)
            return checksum
        if not checksum:
            checksum = make_checksum(fh, algorithm=CHECKSUM_ALGORITHM)
            if self.exists(checksum):
                self.log.debug("Blob already exists, skipping", checksum=checksum)
                return checksum
            fh.seek(0)
        with self._store.open(path.ArchiveKey(checksum).blob, "wb") as out:
            stream(fh, out)
        return checksum

    def put_txt(self, checksum: str, text: str, origin: str = DEFAULT_ORIGIN) -> None:
        """Store extracted text for a file.

        Raises:
            ValueError: If ``checksum`` is not a valid SHA256 hex digest or
                ``origin`` is not a safe path component
                (see `validate_origin`).
        """
        key = path.ArchiveKey(checksum).txt(origin)
        self._txts.put(key, text)

    def get_txt(self, checksum: str, origin: str | None = None) -> str | None:
        """Get extracted text for a file. If `origin`, get by this specific
        extraction, otherwise get the first txt value (no guaranteed order)"""
        if origin:
            key = path.ArchiveKey(checksum).txt(origin)
            return self._txts.get(key)
        for value in self._txts.iterate_values(
            prefix=path.ArchiveKey(checksum), glob="*.txt"
        ):
            return value

    def put_data(self, checksum: str, path: str, data: bytes) -> None:
        """Store raw data at the given path in the checksum's archive directory

        Raises:
            ValueError: If ``checksum`` is not a valid SHA256 hex digest or
                ``path`` is empty, absolute, has a ``..`` segment or is
                reserved (``blob`` or a top-level ``*.txt``).
        """
        self._store.put(self._data_key(checksum, path), data)

    def get_data(self, checksum: str, path: str) -> bytes:
        """Get raw data at the given path in the checksum's archive directory

        Raises:
            ValueError: If ``checksum`` is not a valid SHA256 hex digest or
                ``path`` is empty, absolute, has a ``..`` segment or is
                reserved (``blob`` or a top-level ``*.txt``).
        """
        return self._store.get(self._data_key(checksum, path))

    def _data_key(self, checksum: str, key: str) -> str:
        # out here because `put_data` / `get_data` shadow the `path` module
        relative = validate_relative_uri(key)
        if not relative:  # "." normalises to nothing
            raise ValueError(f"Invalid empty data path: `{key}`")
        # the file's content and what `get_txt` reads as extracted text
        if relative == path.ARCHIVE_BLOB or (
            "/" not in relative and relative.endswith(".txt")
        ):
            raise ValueError(f"Reserved data path: `{key}`")
        return join_relpaths(path.ArchiveKey(checksum), relative)

    def touch(self) -> datetime:
        with self._tags.touch(tag.ARCHIVE_UPDATED) as now:
            return now

class EntityRepository(DatasetHandle):
    """Entities and statements of a dataset: writes go to the journal, reads to
    the parquet store, [`flush`][EntityRepository.flush] moves rows between them.

    Example:
        ```python
        repo = get_entities("my_data")

        with repo.writer(origin="import") as writer:
            writer.add_entity(entity)
        repo.flush()

        for entity in repo.query(Query(C(origin="import"))):
            process(entity)
        ```
    """

    def __init__(
        self,
        dataset: str,
        uri: Uri,
    ) -> None:
        super().__init__(dataset, uri)
        if self._is_api and type(self) is EntityRepository:
            raise RuntimeError(
                "`EntityRepository` cannot run against an http uri directly "
                "– resolve the repository via `get_entities()`"
            )
        self.shards = self._model.shards
        self.compression = self._model.compression
        self._journal = get_journal(dataset)
        self.ENTITIES_JSON = EntitiesArtifact(self).key
        self.EXPORTS_STATEMENTS = StatementsArtifact(self).key

    @cached_property
    def statements(self) -> ParquetStore:
        """The dataset's parquet store, built lazily; raises in api mode."""
        if self._is_api:
            raise RuntimeError(
                f"`{type(self).__name__}.statements` is not available in API mode"
            )
        return ParquetStore(self.uri, self.dataset, self.shards)

    def writer(
        self, origin: str | None = None, role: str | None = None
    ) -> BaseJournalWriter:
        """A bulk journal writer – inserts its tail on success, drops it on error.

        Example:
            ```python
            with repo.writer(origin="import") as writer:
                writer.add_entity(entity)
            ```

        Args:
            origin: Origin of the statements written through it.
            role: Who asserts them, for statements that carry no role.
        """
        return self._journal.writer(origin, role)

    def add(
        self,
        entity: EntityProxy,
        origin: str | None = None,
        fragment: str | None = None,
        role: str | None = None,
    ) -> None:
        """Add a single entity to the journal."""
        self.add_many([entity], origin, fragment, role)

    def add_many(
        self,
        entities: Iterable[EntityProxy],
        origin: str | None = None,
        fragment: str | None = None,
        role: str | None = None,
    ) -> None:
        """Add an entity iterator to the journal."""
        with self.writer(origin, role) as writer:
            for entity in entities:
                writer.add_entity(entity, fragment=fragment)

    def flush(self) -> int:
        """Drain the journal into the parquet store.

        A concurrent flush makes this one a no-op, so ``0`` does not mean the
        journal is empty.

        Returns:
            Number of statements appended.
        """
        with Took() as t:
            self.log.info("Flushing journal ...", journal=mask_uri(self._journal.uri))
            total = self.write_batches(self._journal.flush_batches())
        if total:
            self.log.info(
                "Flushed statements from journal to lake",
                count=total,
                took=t.took,
                journal=mask_uri(self._journal.uri),
            )
        return total

    def write_batches(self, tables: Iterable[pa.Table]) -> int:
        """Append `JOURNAL_SCHEMA` tables to parquet – the journal drain's write
        loop. Each table is durable before the next is pulled.

        Args:
            tables: Stream of packed statement tables.

        Returns:
            Number of rows written.
        """
        store = self.statements  # raises in api mode, also for no tables
        total = 0
        for table in tables:
            if not table.num_rows:
                continue
            store.append(table)
            total += table.num_rows
        return total

    def write_stream(self, tables: Iterable[pa.Table]) -> int:
        """Append `JOURNAL_SCHEMA` tables to parquet as one file per partition –
        the bulk imports' write
        ([`write_stream`][ftm_lakehouse.storage.parquet.ParquetStore.write_stream]).

        Args:
            tables: Stream of packed statement tables.

        Returns:
            Number of rows written.
        """
        return self.statements.write_stream(tables)  # raises in api mode

    def merge(self, force: bool = False) -> None:
        """Flush, then [`merge`][ParquetStore.merge] the parquet store – ``force``
        rewrites clean partitions too."""
        store = self.statements  # raises in api mode, ahead of the flush
        self.flush()
        store.merge(force)

    def shard(self, shards: int) -> None:
        """Flush, then [`shard`][ParquetStore.shard] the parquet store – the
        storage half; [`ShardOperation`][ftm_lakehouse.operation.maintenance.ShardOperation]
        writes the new count to ``config.yml``.

        Args:
            shards: Target shard count; ``<= 1`` means a single shard.
        """
        store = self.statements  # raises in api mode, ahead of the flush
        self.flush()
        store.shard(shards)
        self.shards = shards

    def query_statements_data(self, q: Query | None = None) -> Iterator[StatementDict]:
        """[`query_statements`][EntityRepository.query_statements] as plain dicts –
        the stored columns, timestamps as ISO strings, as the api sends them."""
        for row in self.statements._statement_data(q):
            yield cast(
                StatementDict,
                {
                    k: v.isoformat() if isinstance(v, datetime) else v
                    for k, v in row.items()
                },
            )

    def query(
        self, q: Query | None = None, *, flush_first: bool = False
    ) -> StatementEntities:
        """Query entities from the parquet store.

        Args:
            q: Filters, plus ordering / slicing.
            flush_first: Flush the journal first.
        """
        if flush_first:
            self.flush()
        yield from self.statements.query(q)

    def query_statements(
        self, q: Query | None = None, *, flush_first: bool = False
    ) -> Statements:
        """Query `LakehouseStatement` objects from the parquet store.

        Args:
            q: Filters, plus ordering / slicing.
            flush_first: Flush the journal first.
        """
        if flush_first:
            self.flush()
        yield from self.statements.query_statements(q)

    def get(self, entity_id: str, flush_first: bool = False) -> StatementEntity | None:
        """Get a single entity by ID."""
        q = Query(M(entity_id=entity_id))
        for entity in self.query(q, flush_first=flush_first):
            return entity
        return None

    def stream(self) -> ValueEntities:
        """Stream the exported ``entities.ftm.json`` – not the parquet store."""
        if self._store.exists(self.ENTITIES_JSON):
            with self._store.open(
                self.ENTITIES_JSON, "rb", compression=self.compression
            ) as raw:
                yield from smart_read_proxies(raw)

    def delete_entity(self, entity_id: str, origin: str | None = None) -> int:
        """Tombstone every statement of an entity, parquet and journal – one per
        live row, carrying its ``fragment`` and ``role`` so it shadows that row.

        Args:
            entity_id: The entity to delete.
            origin: Only delete its statements from this origin.

        Returns:
            Number of tombstones written.
        """
        now = utc_now()
        stmts = self._collect_entity_statements(entity_id)
        if origin:
            stmts = [s for s in stmts if s.origin == origin]
        with self.writer() as w:
            for stmt in stmts:
                w.add_statement(stmt, deleted_at=now)
        return len(stmts)

    def delete_statement(
        self,
        stmt: Statement,
        fragment: str | None = None,
        role: str | None = None,
    ) -> None:
        """Tombstone one statement.

        Args:
            stmt: The statement – a `LakehouseStatement` read back from
                [`query_statements`][EntityRepository.query_statements] carries
                its own fragment and role.
            fragment: The row's fragment, for a plain ``Statement``.
            role: The row's role, for a plain ``Statement``.
        """
        with self.writer() as w:
            w.add_statement(stmt, deleted_at=utc_now(), fragment=fragment, role=role)

    def delete_origin(self, origin: str) -> None:
        """Physically drop an origin: flush, then
        [`ParquetStore.delete_origin`][ParquetStore.delete_origin]. Rows journalled
        into it while this runs survive – stop the writers to be sure.

        Args:
            origin: The origin to drop.

        Raises:
            ValueError: If ``origin`` is not a safe origin name.
            RuntimeError: When the write fence cannot be acquired.
        """
        # validate before flushing – a bad origin must cost nothing
        origin = validate_origin(origin)
        self.flush()
        self.statements.delete_origin(origin)

    def _collect_entity_statements(self, entity_id: str) -> list[LakehouseStatement]:
        """An entity's statements from parquet and journal, one per
        ``dedupe_key`` – the journal's win."""
        stmts_by_key: dict[str, LakehouseStatement] = {}

        q = Query(M(entity_id=entity_id))
        for stmt in self.query_statements(q):
            stmt = cast(LakehouseStatement, stmt)
            if stmt.id:
                stmts_by_key[stmt.dedupe_key] = stmt

        for stmt in self._journal.iterate_entity(entity_id):
            if stmt.id:
                stmts_by_key[stmt.dedupe_key] = stmt

        return list(stmts_by_key.values())

    @property
    def version(self) -> int | None:
        """Current version of the main Delta table."""
        return self.statements.version