Layer 4: Operation

class DatasetJobOperation(DatasetHandle, Generic[DJ]):
    """A (long-running) dataset operation, skipped while its target tag is
    newer than its dependencies; the job result is stored after a successful run.

    Repositories come from the LRU-cached factories. Subclasses set `target` and
    `dependencies`, or override `get_target()` / `get_dependencies()`.
    """

    target: str = ""  # tag that gets touched after successful run
    dependencies: list[str] = []  # dependencies for freshness check

    def __init__(self, job: DJ, uri: Uri | None = None) -> None:
        super().__init__(job.dataset, dataset_uri(job.dataset, uri))
        self.job = job
        self.log = job.log

    @cached_property
    def archive(self) -> ArchiveRepository:
        return get_archive(self.dataset, self.uri)

    @cached_property
    def artifacts(self) -> ArtifactsRepository:
        return get_artifacts(self.dataset, self.uri)

    @cached_property
    def entities(self) -> EntityRepository:
        return get_entities(self.dataset, self.uri)

    @cached_property
    def documents(self) -> DocumentRepository:
        return get_documents(self.dataset, self.uri)

    @cached_property
    def jobs(self) -> JobRepository:
        return get_jobs(self.dataset, self.job.__class__, self.uri)

    def get_target(self) -> str:
        """Return the target tag. Override for dynamic values."""
        return self.target

    def get_dependencies(self) -> list[str]:
        """Return the dependencies. Override for dynamic values."""
        return self.dependencies

    def handle(self, run: JobRun, *args, **kwargs) -> None:
        raise NotImplementedError

    def prepare(self) -> None:
        """Bring the dataset into the state `handle` reads from – no-op by default.

        Runs before the freshness check and outside the target's window:
        `Tags.touch` stamps the target with its *entry* time, so a dependency
        moved inside the window would mark the result stale at once.
        [`ExportOperation`][ftm_lakehouse.operation.export.ExportOperation]
        drains the journal here.
        """

    def is_fresh(self) -> bool:
        """Whether the target tag is newer than every dependency – nothing to
        do. Override where freshness is not a pair of timestamps."""
        target = self.get_target()
        dependencies = self.get_dependencies()
        if not (target and dependencies):
            return False
        return self._tags.is_latest(target, dependencies)

    def _run_local(self, force: bool | None = False, *args, **kwargs) -> DJ:
        """Core run logic – prepare() + orchestration + handle()."""
        target = self.get_target()
        dependencies = self.get_dependencies()

        self.prepare()

        if not force:
            if self.is_fresh():
                self.job.log.info(
                    f"Already up-to-date: `{target}`, skipping ...",
                    target=target,
                    dependencies=dependencies,
                )
                self.job.stop()
                return self.job

        # target tag and job result are stored on a clean exit
        with self.jobs.run(self.job) as run, self._tags.touch(target) as now:
            self.job.log.info(
                f"Start `{target}` ...",
                target=target,
                dependencies=dependencies,
                started=now,
            )
            _ = self.handle(run, *args, force=force, **kwargs)
        self.log.info(
            f"Done `{target}`.",
            target=target,
            dependencies=dependencies,
            started=now,
            took=run.job.took,
            errors=run.job.errors,
        )
        return run.job

    def run(self, force: bool | None = False, *args, **kwargs) -> DJ:
        """Run the operation unless fresh (``force`` runs it anyway); in api mode
        the job is delegated to the server's operations endpoint."""
        if self._is_api:
            return self._api_run(force, *args, **kwargs)
        return self._run_local(force, *args, **kwargs)

    def _api_run(self, force: bool | None = False, *args, **kwargs) -> DJ:
        """Delegate run to remote api"""
        url = self._api.make_url("_api/operations")
        res = self._api.make_request(
            url,
            "POST",
            params={"force": force},
            json=self.job.model_dump(mode="json"),
        )
        return self.job.__class__(**res.json())

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}({self.job.dataset})>"

class CrawlJob(DatasetJobModel):
    """
    Job model for crawl operations.

    Tracks the state and configuration of a crawl job.

    Attributes:
        uri: Source location URI to crawl
        prefix: Include only keys with this prefix
        exclude_prefix: Exclude keys with this prefix
        glob: Include only keys matching this glob pattern
        exclude_glob: Exclude keys matching this glob pattern
        make_entities: Add document entities to statement store
        existing: Skip files whose checksum is already archived, or overwrite
    """

    uri: Uri
    prefix: str | None = None
    exclude_prefix: str | None = None
    glob: str | None = None
    exclude_glob: str | None = None
    make_entities: bool = False
    existing: HandleExistingMode | None = HandleExistingMode.overwrite

class CrawlOperation(DatasetJobOperation[CrawlJob]):
    """
    Crawl workflow that archives files and creates entities.

    Iterates through files in a source store, archives them to the
    file repository, and creates corresponding entities in the
    entities repository.

    Example:
        ```python
        from ftm_lakehouse.operation import CrawlOperation, CrawlJob

        job = CrawlJob.make(
            uri="s3://bucket/documents",
            dataset="my_dataset",
            glob="*.pdf"
        )
        op = CrawlOperation(job=job)
        result = op.run()
        print(f"Crawled {result.done} files")
        ```
    """

    target = tag.OP_CRAWL

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.source = get_store(self.job.uri)
        if self.source.is_http:
            try:
                import aiohttp

                backend_config = ensure_dict(self.source.backend_config)
                backend_config["client_kwargs"] = {
                    **ensure_dict(backend_config.get("client_kwargs")),
                    "timeout": aiohttp.ClientTimeout(total=3600 * 24),
                }
                self.source.backend_config = backend_config
            except ImportError as e:
                raise ImportError(f"Please install `aiohttp` dependency: {e}")

    def get_uris(self) -> Generator[str, None, None]:
        """
        Generate file uris to crawl.

        Applies prefix, glob, and exclude filters to the source store.

        Yields:
            File uris to be crawled
        """
        self.log.info(f"Crawling `{mask_uri(self.job.uri)}` ...")
        for key in self.source.iterate_keys(
            prefix=self.job.prefix,
            exclude_prefix=self.job.exclude_prefix,
            glob=self.job.glob,
        ):
            if self.job.exclude_glob and fnmatch(key, self.job.exclude_glob):
                continue
            self.job.pending += 1
            self.job.touch()
            yield key

    def handle_crawl(self, uri: str, run: JobRun[CrawlJob]) -> datetime:
        """
        Handle a single crawl task.

        Archives the file and creates a corresponding entity.

        Args:
            uri: File uri to crawl
            run: Current job run context

        Returns:
            Timestamp when the task was processed
        """
        now = utc_now()

        self.log.info(f"Crawling `{uri}` ...", source=mask_uri(self.source.uri))
        checksum = None
        if self.source.is_local:
            checksum = self.source.checksum(uri, algorithm=CHECKSUM_ALGORITHM)
        if not self._should_skip(checksum):
            checksum = self.archive.store(
                self.source.to_uri(uri), checksum=checksum, tag_updated=False
            )
            if self.job.make_entities:
                info = self.source.info(uri)
                entities = make_file_entities(info, checksum, self.dataset)
                self.entities.add_many(entities, tag.CRAWL_ORIGIN)
            run.job.done += 1
        return now

    def handle(self, run: JobRun, *args, **kwargs) -> None:
        for ix, task in enumerate(self.get_uris(), 1):
            if ix % 1000 == 0:
                self.log.info(
                    f"Handling task {ix} ...",
                    pending=self.job.pending,
                    done=self.job.done,
                )
                run.save()
                self.archive._tags.set(tag.ARCHIVE_UPDATED)
            self.handle_crawl(task, run)
            run.job.pending -= 1
            run.job.touch()
        self.archive._tags.set(tag.ARCHIVE_UPDATED)
        if self.job.make_entities:
            self.entities.flush()

    def _should_skip(self, checksum: str | None) -> bool:
        if self.job.existing != HandleExistingMode.skip_checksum or checksum is None:
            return False
        return self.archive.exists(checksum)

    def _api_run(self, force: bool | None = False, *args, **kwargs) -> CrawlJob:
        """Crawl always runs locally – source files aren't on the API server."""
        return self._run_local(force, *args, **kwargs)

class ExportKind(StrEnum):
    """What each export artifact is called (`Artifact.kind`) – also the key an
    export run reports its counts under."""

    statements = "statements"
    entities = "entities"
    documents = "documents"
    parents = "parents"
    statistics = "statistics"
    index = "index"  # type: ignore[assignment]  # shadows str.index, fine for enums

class ExportJob(DatasetJobModel):
    """Job model for the export."""

    make_diff: bool = True
    """Also export the delta diff files of the diffable artifacts."""
    result: dict[str, int] | None = None
    """What the run wrote, per artifact and per diff op."""

class ExportOperation(DatasetJobOperation[ExportJob]):
    """Export the dataset in one sweep over the entity stream – skipped while
    the exports are newer than the last content change."""

    target = tag.OP_EXPORT
    dependencies = [tag.STATEMENTS_UPDATED]
    """The content clock – a merge does not move it."""

    def prepare(self) -> None:
        """Drain the journal ahead of the freshness check, so buffered rows
        count."""
        self.entities.flush()

    def export(self, now: datetime) -> dict[str, int]:
        """Write every streamed artifact from one pass over the entities.

        Each ``(shard, bucket)`` pair is swept into parts of every artifact in
        ``LAKEHOUSE_WORKERS`` processes; an artifact is replaced only once every
        pair is in, so a failed run leaves the previous one intact – and the
        diff states and ``statistics.json`` only once every artifact is. Held
        under the merge lock, so an ``optimize`` cannot vacuum the snapshot's
        files.

        Args:
            now: When the run started – diff files are named after it.

        Returns:
            Counts per artifact and per diff op.
        """
        workers = settings.workers
        store = self.entities.statements
        counts: Counter[str] = Counter()
        with (
            store.merge_lock(),
            TemporaryDirectory(prefix="ftm-lakehouse-export-") as tmp,
        ):
            # no diff asked for is no version: no diff window, no state recorded
            version = store.version if self.job.make_diff else None
            sources = store.sweep_sources()
            # the parent's own part: the documents csv and the DELs of `finish`
            session = ExportSession(
                self.artifacts.runs(now, f"{tmp}/parent"),
                version,
                store.deleted_candidates,
            )
            parts = [f"{tmp}/{s.key[0]}-{s.key[1]}" for s in sources]
            streamed = list(self.artifacts.streamed())
            # one file per artifact and diff series, each part appended as it comes
            assemblies: list[tuple[Callable[[str], str], Assembly]] = [
                *(
                    (a.part, a.assembly())
                    for a in streamed
                    if not isinstance(a, VersionedArtifact)
                ),
                *(
                    (a.diff_part, a.diff_assembly(now))
                    for a in streamed
                    if isinstance(a, DiffableArtifact)
                ),
            ]

            def append(parts: str) -> None:
                for part, assembly in assemblies:
                    assembly.append(part(parts))

            try:
                with (
                    session,
                    SyncProgressBar("Exporting...", len(sources)) as bar,
                    process_map(workers) as run,
                ):
                    pending = self._pending_by_shard(session)
                    tasks = [
                        ExportTask(
                            dataset=self.dataset,
                            uri=str(self.uri),
                            now=now,
                            version=version,
                            parts=part,
                            source=source,
                            duckdb_config=worker_duckdb_config(workers),
                            pending=pending.get(source.key[0], {}),
                        )
                        for part, source in zip(parts, sources)
                    ]
                    # largest first, so no big pair starts last and holds up the end
                    tasks.sort(key=lambda task: task.source.size, reverse=True)
                    by_parts = {task.parts: task for task in tasks}
                    headed = False
                    for part in run(export_partition, tasks):
                        shard, bucket = by_parts[part.parts].source.key
                        statements = part.counts.get("statements", 0)
                        counts.update(part.counts)
                        if statements and not headed:
                            # ahead of the first row, so an empty csv stays empty
                            append(self._write_header(tmp))
                            headed = True
                        append(part.parts)
                        session.adopt(part.parts, part.seen, part.stats, part.counts)
                        # the bar's throughput: the bytes each pair read
                        bar.advance(size=by_parts[part.parts].source.size)
                        self.log.info(
                            f"Processed `{shard}/{bucket}`.",
                            took=part.took,
                            shard=shard,
                            bucket=bucket,
                            statements=statements,
                        )
                    self.log.info(
                        "Swept every pair, writing the documents ...",
                        pairs=len(sources),
                    )
                # after the session closed: its own writers' codec trailers are
                # written, so its part is a complete frame
                with Took() as t:
                    append(f"{tmp}/parent")
                    for _, assembly in assemblies:
                        assembly.commit()
                self.log.info("Moved the artifacts into place.", took=t.took)
                session.commit()
            except BaseException:
                for _, assembly in assemblies:
                    assembly.abort()
                raise
            counts.update(session.result())
        return dict(counts)

    def _write_header(self, tmp: str) -> str:
        """The ``statements.csv`` header as a part directory of its own."""
        artifact = self.artifacts.statements
        part = f"{tmp}/header"
        Path(part).mkdir(parents=True, exist_ok=True)
        with smart_open(
            artifact.part(part), "wb", compression=artifact.compression
        ) as fh:
            fh.write(statement_csv_header())
        return part

    def _pending_by_shard(
        self, session: ExportSession
    ) -> dict[str, dict[str, frozenset[str]]]:
        """Each diff series' DEL candidates per shard – not per pair: an id
        names its shard but not its bucket."""
        shards = self.entities.shards
        out: dict[str, dict[str, set[str]]] = {}
        for run in session.diffable:
            for entity_id in run.pending:
                shard = entity_shard(entity_id, shards)
                out.setdefault(shard, {}).setdefault(run.name, set()).add(entity_id)
        return {
            shard: {name: frozenset(ids) for name, ids in runs.items()}
            for shard, runs in out.items()
        }

    def export_index(self) -> None:
        """Write ``index.json``, registering what the exports produced."""
        dataset = self._model
        dataset.resources = list(self.artifacts.resources())
        statistics = self.artifacts.statistics
        if statistics.exists():
            dataset.apply_stats(self._store.get(statistics.key, model=DatasetStats))
        self.artifacts.index.write(dataset)

    def handle(self, run: JobRun[ExportJob], *args: Any, **kwargs: Any) -> None:
        """The sweep, then every artifact's freshness tag – none if it crashed –
        then ``index.json``, also for an empty store."""
        if self.entities.statements.exists:
            started = utc_now()
            result = self.export(started)
            for artifact in self.artifacts.streamed():
                artifact.touch(started)
            self.log.info("Export(s) done.", **result)
            run.job.result = result
        else:
            self.log.info(
                "Statement store empty, nothing to sweep ...",
                uri=mask_uri(self.entities.uri),
            )
        self.export_index()
        run.job.done = 1

class DiffOp(StrEnum):
    """What a diff entry says happened to an entity.

    Ref. https://www.opensanctions.org/docs/bulk/delta/
    """

    ADD = "ADD"
    """The entity is new – every statement it has arrived in this window."""

    MOD = "MOD"
    """The entity predates this window and changed in it – it gained
    statements, or lost some to a tombstone while staying alive."""

    DEL = "DEL"
    """The entity is gone entirely."""

class OptimizeJob(DatasetJobModel):
    retention_hours: int = 0
    """Vacuum: retain obsolete files newer than this many hours."""

class OptimizeOperation(DatasetJobOperation[OptimizeJob]):
    """Optimize the parquet statement store: merge, then vacuum.

    Merge rewrites each dirty ``(shard, bucket, origin)`` partition into one
    file, keeping the latest row per statement id with ``min(first_seen)`` and
    ``max(last_seen)``. Tombstones past the grace period are dropped; vacuum
    deletes the replaced files.
    """

    target = tag.OP_OPTIMIZE
    dependencies: list[str] = []

    def is_fresh(self) -> bool:
        return not self.entities.statements.needs_merge

    def handle(self, run: JobRun[OptimizeJob], force: bool = False, **kwargs) -> None:
        self.entities.merge(force)
        run.job.done += 1
        run.save()
        self.entities.statements.vacuum(retention_hours=run.job.retention_hours)
        run.job.done += 1

class ShardJob(DatasetJobModel):
    shards: int = Field(ge=0)
    """Target entity-id hash shard count; ``0`` / ``1`` means a single shard."""

class ShardOperation(DatasetJobOperation[ShardJob]):
    """Change the dataset's shard count: rewrite the store, then the config.

    The rewrite drains the journal and moves every ``(bucket, origin)`` group
    into the new shard partitions (one atomic Delta commit per group); then the
    count is written to ``config.yml`` and the factory caches are cleared. The
    config goes last – it declares the layout to every other process – so a run
    that dies in between is repaired by running it again (the rewrite is
    idempotent). Every partition comes out dirty: run ``optimize`` afterwards.
    """

    target = tag.OP_SHARD

    def is_fresh(self) -> bool:
        """Whether the config already names the target count – not a tag pair,
        so a hand-edited config reads as fresh; ``force`` overrides."""
        return self._model.shards == self.job.shards

    def handle(self, run: JobRun[ShardJob], **kwargs: Any) -> None:
        self.entities.shard(self.job.shards)
        run.job.done += 1
        run.save()
        self._versions.make(
            path.CONFIG, self._model.model_copy(update={"shards": self.job.shards})
        )
        factories.clear_caches()
        run.job.done += 1