Layer 4: Operation

class DatasetJobOperation(DatasetHandle, Generic[DJ]):
    """
    A (long-running) operation for a specific dataset that updates tags and
    checks dependencies for freshness to be able to skip this operation. The job
    result is stored after successful run.

    Repositories are resolved through the LRU-cached factories, so an
    operation shares its repository instances with every other path that
    addresses the same dataset.

    Subclasses can either set class attributes `target` and `dependencies`,
    or override `get_target()` and `get_dependencies()` for dynamic values.
    """

    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.

        Runs *before* the freshness check and before the target tag's window
        opens, which is what makes it usable at all: `Tags.touch` stamps
        the target with the timestamp it *entered*, so preparation that writes
        a dependency tag from inside the window would mark the result stale the
        moment it is written. Ahead of the window, the timestamps stay honest –
        prepare moves the dependency, then the target is stamped after it.

        No-op by default;
        [`ExportOperation`][ftm_lakehouse.operation.export.ExportOperation] drains the
        journal and merges the statement store here, since exports read
        canonical rows.
        """

    def is_fresh(self) -> bool:
        """Whether the target is newer than every dependency – nothing to do.

        Tag-pair comparison by default. Override where the question is not a
        pair of timestamps (``OptimizeOperation`` asks the statement store
        directly).
        """
        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

        # Execute: Store target tag and job result on successful context leave
        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:
        """Execute the handle function, force to run it regardless of freshness
        dependencies. In api mode the whole job is delegated to the remote
        operations endpoint (`_api_run`)."""
        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
        store_metadata: Write file.json metadata alongside archive blobs
    """

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

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(uri, checksum):
            file = self.archive.store(
                self.source.to_uri(uri),
                checksum=checksum,
                store_metadata=self.job.store_metadata,
                tag_updated=False,
                key=uri,
                origin=tag.CRAWL_ORIGIN,
            )
            if self.job.make_entities:
                self.entities.add_many(file.make_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, uri: Uri, checksum: str | None) -> bool:
        if self.job.existing is None:
            return False
        if self.job.existing == HandleExistingMode.overwrite:
            return False
        if checksum is None:
            return False
        if (
            self.job.existing == HandleExistingMode.skip_checksum
            or not self.job.store_metadata
        ):
            return self.archive.exists(checksum)
        if self.job.existing == HandleExistingMode.skip_path:
            if self.archive.exists(checksum):
                for file in self.archive.get_all_files(checksum):
                    if file.key == str(uri):
                        return True
        return False

    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):
    """The available dataset exports.

    Lives here rather than with the export operation because each `Artifact`
    declares the kind it answers to, and ``repository/`` cannot import
    ``operation/``. `ftm_lakehouse.operation.export` re-exports it.
    """

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

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

    kind: ExportKind = ExportKind.all
    make_diff: bool = True
    """Also export delta diff files (``entities`` / ``documents`` kinds)."""
    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.

    Flushes and merges first ([`prepare`][ExportOperation.prepare]) – exports
    read canonical rows. Skips if the target is newer than the last optimize.

    A run stamps a freshness tag per artifact it wrote, so a later single-kind
    export sees itself up to date and ``index.json`` still finds the
    dependencies it registers.
    """

    @cached_property
    def kinds(self) -> tuple[ExportKind, ...]:
        """The sweep artifacts this run writes."""
        if self.job.kind == ExportKind.all:
            return SWEEP_KINDS
        if self.job.kind in SWEEP_KINDS:
            return (self.job.kind,)
        return ()

    def get_target(self) -> str:
        if self.job.kind == ExportKind.all:
            return tag.OP_EXPORT
        return str(self.artifacts[self.job.kind].tag)

    def get_dependencies(self) -> list[str]:
        if self.job.kind == ExportKind.all:
            return [tag.STATEMENTS_OPTIMIZED]
        return [str(d) for d in self.artifacts[self.job.kind].dependencies]

    def prepare(self) -> None:
        """Drain the journal and merge, so the export reads canonical rows."""
        if not self._tags.is_latest(tag.JOURNAL_FLUSHED, [tag.JOURNAL_UPDATED]):
            self.entities.flush()
        if self.entities.exists and self.entities.needs_merge:
            self.entities.merge()

    def iterate(self) -> Iterator[EntityPayload]:
        """Every entity in the store, folded from one scan.

        Writes ``statements.csv`` from the same Arrow batches when this run
        covers it, so the csv costs a tee rather than a second pass. Rows are
        only materialised when something downstream needs them – a
        statements-only export stays columnar end to end.
        """
        with_csv_export = ExportKind.statements in self.kinds
        tee = bool({ExportKind.entities, ExportKind.documents} & set(self.kinds))
        rows = self.entities.sweep(with_csv_export=with_csv_export, tee=tee)
        yield from aggregate_unsafe(rows, self.dataset)

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

        Args:
            now: Timestamp the run started – the diff files are named after it
                and the diff states are recorded at it.

        Returns:
            Counts per artifact and per diff op.
        """
        version = self.entities.version
        if self.job.make_diff and version is not None and self.entities.needs_merge:
            raise RuntimeError(
                "Cannot export diffs: the statement store has un-merged writes "
                "and a diff publishes canonical entities. Run "
                "`ftm-lakehouse maintenance optimize` first."
            )
        session = self.artifacts.session(now, self.kinds, version, self.job.make_diff)
        with session:
            for payload in self.iterate():
                session.consume(payload)
        return session.result()

    def export_statistics(self) -> None:
        """Write ``statistics.json`` from the store's global SQL aggregate."""
        self.artifacts.statistics.write(self.entities.stats())

    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:
        if run.job.kind == ExportKind.index:
            self.export_index()
            run.job.done = 1
            return

        if not self.entities.exists:
            self.log.info(
                "Statement store empty, skipping ...",
                uri=mask_uri(self.entities.uri),
            )
            return

        if run.job.kind == ExportKind.statistics:
            self.export_statistics()
            run.job.done = 1
            return

        started = utc_now()
        result = self.export(started)
        for artifact in self.artifacts.written_by(self.kinds):
            artifact.touch(started)
        self.log.info("Export(s) done.", **result)
        run.job.result = result
        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, compact, vacuum.

    For each ``(shard, bucket, origin)`` partition: keep the most-recent row
    per statement id, fold ``first_seen`` down to the minimum, drop tombstones
    older than the grace period – then bin-pack small files and delete
    obsolete ones. Each step is held under the dataset write fence.
    """

    target = tag.STATEMENTS_OPTIMIZED
    dependencies = [tag.STATEMENTS_UPDATED]

    def is_fresh(self) -> bool:
        """Ask the statement store whether any partition is unmerged.

        The tag pair cannot answer this one. ``merge`` stamps
        [`STATEMENTS_OPTIMIZED`][ftm_lakehouse.core.conventions.tag.STATEMENTS_OPTIMIZED] on
        completion while the target tag records when this operation *started*,
        so a successful optimize always finishes behind its own dependency and
        reads as stale – costing a redundant full pass every time. The
        per-partition tags
        [`ParquetStore.merge`][ftm_lakehouse.storage.parquet.ParquetStore.merge]
        compares internally are the sound predicate, and ``needs_merge`` is
        that comparison.
        """
        return not self.entities.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.compact()
        run.job.done += 1
        run.save()
        self.entities.vacuum(retention_hours=run.job.retention_hours)
        run.job.done += 1

class ShardJob(DatasetJobModel):
    shards: int = Field(ge=0)
    """Target number of entity-id hash shards. ``0`` / ``1`` means a single
    shard; the value is bounded below because it becomes a partition key."""

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

    The shard count is otherwise fixed at creation – every reader and
    writer resolves it from ``config.yml`` – so growing it is a full
    rewrite of the statement store. The typical trigger is a dataset that
    outgrew its layout: one shard means one partition per
    ``(bucket, origin)``, and queries that have to scan it whole get
    slow.

    Two steps, in this order:

    1. `shard`
       drains the journal and rewrites every ``(bucket, origin)`` group
       into the new shard partitions, streamed, one atomic Delta commit
       per group.
    2. the new count is written to ``config.yml`` (versioned like every
       other config write) and the repository factory caches are
       invalidated, so repositories fetched afterwards resolve the new
       layout.

    The config write goes last on purpose: it is what declares the layout
    to every other process, so it must not run ahead of the data. A run
    that dies in between leaves the config on the old count and is
    repaired by running it again – the rewrite recomputes each shard from
    ``entity_id`` alone, so it is idempotent.

    The rewrite is neither sorted nor deduped, which leaves every
    partition marked dirty – run ``optimize`` afterwards to restore
    canonical content and file sort order.
    """

    target = tag.OP_SHARD

    def is_fresh(self) -> bool:
        """Whether the dataset is already configured for the target count.

        Not a tag pair: what a re-shard changes is the configured layout,
        so the config *is* the freshness state. Consequently a config
        edited by hand to a count the store was never rewritten for reads
        as fresh – ``force`` is the way out of that.
        """
        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

class MigrateJob(DatasetJobModel):
    """No parameters – a migrate run is always "everything outstanding"."""

class MigrateOperation(DatasetJobOperation[MigrateJob]):
    """Apply the storage-layout migrations this dataset has not seen yet.

    Runs the functions registered in ``ftm_lakehouse.operation.migrations`` in
    registry order, stamping each with
    [`tag.migration`][ftm_lakehouse.core.conventions.tag.migration] on
    completion. Per-migration tags rather than one version number: a run that
    dies halfway keeps what it finished and the next one picks up at the first
    untagged migration. ``force`` re-runs the whole registry – migrations are
    idempotent.
    """

    target = tag.OP_MIGRATE

    @property
    def outstanding(self) -> tuple[Migration, ...]:
        """The registered migrations this dataset carries no tag for."""
        return tuple(
            m for m in MIGRATIONS if self._tags.get(tag.migration(m.__name__)) is None
        )

    def is_fresh(self) -> bool:
        """Whether every registered migration has run against this dataset.

        Not a tag pair: a migration is done or not, and no dependency's
        timestamp can make an applied one stale again.
        """
        return not self.outstanding

    def handle(
        self, run: JobRun[MigrateJob], force: bool = False, **kwargs: Any
    ) -> None:
        migrations = MIGRATIONS if force else self.outstanding
        ref = DatasetRef(self.dataset, str(self.uri))
        run.job.pending = len(migrations)
        run.save()
        for migration in migrations:
            name = migration.__name__
            with self._tags.touch(tag.migration(name)), Took() as t:
                self.log.info(f"Running migration `{name}` ...", migration=name)
                migration(ref)
                run.job.pending -= 1
                run.job.done += 1
                run.save()
                self.log.info(f"Migration `{name}` done", migration=name, took=t.took)

class MakeJob(DatasetJobModel):
    pass

class MakeOperation(DatasetJobOperation[MakeJob]):
    target = tag.OP_MAKE
    dependencies = [tag.JOURNAL_UPDATED, tag.STATEMENTS_OPTIMIZED]

    def prepare(self) -> None:
        """Drain the journal; each export merges for itself in its own
        [`ExportOperation.prepare`][ExportOperation.prepare]."""
        self.entities.flush()

    def handle(self, run: JobRun, *args, **kwargs) -> None:
        """Run the export sweep, then the two artifacts computed from it."""
        force = kwargs.get("force", False)
        for kind in MAKE_KINDS:
            job = ExportJob.make(dataset=self.dataset, kind=kind)
            ExportOperation(job, self.uri).run(force=force)
        run.job.done = 1