oarepo_oaipmh_harvester.writers

class OAIServiceWriter(BaseWriter):
    """Writer for OAI-PMH harvested records.

    This writer uses a service to write records harvested via OAI-PMH. It also stores
    an OAIRecord model instance to keep track of the harvesting status.
    """

    def __init__(  # noqa PLR0913,PLR0917
        self,
        model: str,
        *args: Any,
        identity: Identity | None = None,
        update_all: bool = False,
        harvester_id: str | None = None,
        pid_field: str = "id",
        publish: bool = True,
        direct_update: bool = False,
        community: str | None = None,
        **kwargs: Any,
    ):
        """Initialize the OAI service writer.

        Args:
            model: Name of the target model (`current_runtime.models`).
            identity: Identity used for the service calls, defaults to the system identity.
            update_all: Write every record, even if it is up to date.
            harvester_id: Harvester the harvest records belong to (set by the harvest task).
            pid_field: Path of the pid value in the service result.
            publish: Publish created / updated drafts. Required (`ValueError` otherwise) for
                rdm models; `False` only for models without drafts.
            direct_update: Update published records via `service.update` instead of edit -
                update draft - publish. Only for services that allow updating published
                records (drafts services do not by default).
            community: Slug (or id) of a community every written record is added to. Only
                for rdm models.

        """
        self._model = model
        self._update_all = update_all
        self._harvester_id = harvester_id
        self._pid_field = pid_field
        self._identity = identity or system_identity
        self._publish = publish
        self._direct_update = direct_update
        self._community = community

        super().__init__(
            *args,
            **kwargs,
        )

        is_rdm = hasattr(self.service, "delete_record")
        if not publish and is_rdm:
            # rdm records can not be updated directly, only via edit - update draft - publish
            raise ValueError(f"Model '{model}' uses rdm records service, the oai-service writer needs publish=True.")
        if community and not is_rdm:
            raise ValueError(f"Model '{model}' does not use rdm records service, it can not be added to communities.")

    @property
    def service(self) -> RecordService:
        """Get the service for the specified model."""
        return current_runtime.models[self._model].service

    @override
    def write(  # too complex
        self,
        stream_entry: StreamEntry,
        *args: Any,
        **kwargs: Any,
    ) -> StreamEntry:
        """Write the input entry using a given service."""
        current_app.logger.debug("Writing entry: %s", stream_entry.entry)
        harvested_at = dt.datetime.now(dt.UTC)
        original_data = {"oai_payload": stream_entry.entry["oai_record"].raw}
        transformed_data = stream_entry.entry["record"]

        # 0. extract the record oai identifier
        oai_identifier = stream_entry.entry["oai_record"].header.identifier
        oai_datestamp = _utc_datestamp(stream_entry.entry["oai_record"].header.datestamp)
        oai_deleted = stream_entry.entry["oai_record"].header.deleted

        # 1. check if there is an OAI record already
        oai_record = self._get_oai_record(oai_identifier)

        current_app.logger.debug(
            "Processing OAI identifier %s: existing OAI record: %s",
            oai_identifier,
            oai_record,
        )

        written_data = None
        validation_errors: list[dict] = []
        exception_raised = None
        op_type = None

        try:
            if stream_entry.errors or stream_entry.exc is not None:
                # reading / transformation failed - only the errors are stored on the oai
                # record, the local record is left as it is
                current_app.logger.debug("Not writing OAI identifier %s, it has errors", oai_identifier)
            else:
                pid_value = oai_record.record_pid if oai_record else None
                if not pid_value and not oai_deleted:
                    pid_value = self._get_pid_value_from_record(transformed_data)

                # 2. find out the state of the local record
                local = LocalRecord.resolve(self.service, pid_value, with_drafts=self._publish)

                # 3. decide on operation type
                op_type = local.operation(oai_deleted)

                # 4. if exists, check if the record is already up to date
                if self._check_up_to_date(stream_entry, oai_record, oai_datestamp, oai_deleted, local):
                    return stream_entry
                current_app.logger.debug(
                    "Determined operation type '%s' for OAI identifier %s (pid_value=%s, local state=%s)",
                    op_type,
                    oai_identifier,
                    pid_value,
                    local.state,
                )

                written_data, validation_errors = self._save_and_publish_record(
                    transformed_data,
                    local,
                    op_type,
                    oai_identifier,
                )
        except Exception as e:  # noqa BLE001 ok to catch here
            # we can't know the state of the db connection, rollback is a safe bet
            db.session.rollback()
            exception_raised = e

        # In case of exception, the session has been rolled back in the lines
        # above, so we can safely proceed to store the OAI record
        stored_errors = self._store_oai_record(
            oai_record,
            oai_identifier,
            oai_datestamp,
            oai_deleted,
            harvested_at,
            original_data,
            transformed_data,
            stream_entry,
            validation_errors,
            exception_raised,
            written_data,
        )
        if stored_errors:
            current_app.logger.info(
                "OAI identifier %s: %s error(s) stored on the harvest record: %s",
                oai_identifier,
                len(stored_errors),
                "; ".join(_short_error(error) for error in stored_errors),
            )
        if exception_raised is not None and stream_entry.exc is None:
            # report the failure to the datastream as well, not only on the oai record
            stream_entry.exc = exception_raised

        # the datastream api expects error messages; structured errors are on the oai record
        stream_entry.errors = [
            *(stream_entry.errors or []),
            *(f"{error.get('location', '')}: {error['message']}" for error in validation_errors),
        ]
        stream_entry.entry = {
            "record": written_data,
            "oai_record": stream_entry.entry["oai_record"],
        }
        stream_entry.op_type = op_type
        return stream_entry

    def _save_and_publish_record(
        self,
        transformed_data: dict,
        local: LocalRecord,
        op_type: Operation,
        oai_identifier: str,
    ) -> tuple[dict | None, list[dict]]:
        """Save and optionally publish a record within a unit of work.

        Args:
            transformed_data: The transformed record data to be saved.
            local: The local record (state, pid value) the operation is applied to.
            op_type: The operation type (create, update, delete, etc.).
            oai_identifier: The OAI identifier of the record.

        Returns:
            written_data: The data of the written record, or None if not applicable.
            validation_errors: Errors reported by the service (`type`, `location`,
                `message`); the transaction is rolled back if there are any.

        """
        # we use unit of work here to avoid side effects
        # (such as created draft record) in case of failure
        with UnitOfWork() as uow:
            # resolve lazy strings before writing
            transformed_data = json.loads(json.dumps(transformed_data, default=str))
            # parent communities are not writable via the service, they are added after publishing
            communities = self._pop_communities(transformed_data)
            self._add_files_defaults(transformed_data)
            try:
                written_data = self._write(transformed_data, local, op_type, oai_identifier, uow)
            except Exception:
                current_app.logger.exception(
                    "Error during '%s' operation for OAI identifier %s",
                    op_type,
                    oai_identifier,
                )
                raise
            commit_transaction = True
            validation_errors: list[dict] = []
            if written_data:
                # ensure the written data is also transformed (e.g. to resolve lazy strings)
                written_data = json.loads(json.dumps(written_data, default=str))

                current_app.logger.debug(
                    "Written (%s) %s",
                    op_type,
                    json.dumps(written_data, indent=2, ensure_ascii=False),
                )

                if written_data.get("errors"):
                    # the service reported errors during creation/update

                    validation_errors = [x for err in written_data["errors"] for x in self._convert_to_oai_error(err)]
                    commit_transaction = False
                elif self._publish and (
                    op_type in (Operation.CREATE, Operation.UPDATE_DRAFT)
                    or (
                        # with direct update the published record is updated, no draft to publish
                        not self._direct_update and op_type in (Operation.UPDATE, Operation.RESTORE, Operation.REVIVE)
                    )
                ):
                    draft_service = cast("DraftRecordService", self.service)
                    draft_service.publish(
                        self._identity,
                        dict_lookup(written_data, self._pid_field),
                        uow=uow,
                    )
                if commit_transaction:
                    self._add_to_communities(dict_lookup(written_data, self._pid_field), communities, uow)
            if commit_transaction:
                uow.commit()
            else:
                uow.rollback()
            return written_data, validation_errors

    def _add_files_defaults(self, record_data: dict | None) -> None:
        """Disable files / media files of records whose data do not mention them.

        Only for the sections the model supports - a model without media files would reject
        the `media_files` key.
        """
        if record_data is None:
            # deleted at the source
            return
        model = current_runtime.models[self._model]
        if model.file_service is not None:
            record_data.setdefault("files", {"enabled": False})
        if model.media_file_service is not None:
            record_data.setdefault("media_files", {"enabled": False})

    def _pop_communities(self, record_data: dict | None) -> list[str]:
        """Remove `parent.communities` from the record data, return the writer's and the listed communities."""
        # no record data for records deleted at the source
        parent = (record_data or {}).get("parent") or {}
        listed = (parent.pop("communities", None) or {}).get("ids") or []
        # dict.fromkeys - remove duplicates, keep the order
        return list(dict.fromkeys([*([self._community] if self._community else []), *listed]))

    def _add_to_communities(self, pid_value: str, communities: list[str], uow: UnitOfWork) -> None:
        """Add the published record to the communities (slugs or ids) it is not in yet.

        `bulk_add` adds the record directly, without an inclusion request (allowed for system
        processes only), also adds a subcommunity's parent and makes the community the default
        one when the record has none.
        """
        for community in communities:
            # resolved for each community - adding a subcommunity also adds its parent
            record = cast("RDMRecord", self.service.record_cls.pid.resolve(pid_value))
            if current_communities.service.record_cls.pid.resolve(community).id in record.parent.communities:
                continue
            errors = current_rdm_records.record_communities_service.bulk_add(
                self._identity, community, [pid_value], uow=uow
            )
            if errors:
                raise ValueError(f"Can not add record {pid_value} to community {community}: {errors[0]['message']}")

    def _get_oai_record(self, oai_identifier: str) -> OAIHarvestedRecord | None:
        """Get the OAIHarvestedRecord instance by its OAI identifier."""
        return db.session.query(OAIHarvestedRecord).filter_by(oai_identifier=oai_identifier).one_or_none()

    def _check_up_to_date(
        self,
        stream_entry: StreamEntry,
        oai_record: OAIHarvestedRecord | None,
        oai_datestamp: dt.datetime,
        oai_deleted: bool,
        local: LocalRecord,
    ) -> bool:
        """Check if the local record already reflects the harvested one, so writing can be skipped.

        Besides unchanged oai metadata, the local record must be in the harvested state -
        a record deleted or restored locally since the last harvest is written again.
        """
        if not oai_record:
            return False

        if (
            oai_datestamp == oai_record.datestamp
            and oai_deleted == oai_record.deleted
            and not self._update_all
            and oai_record.has_errors is False
            and oai_record.has_warnings is False
        ):
            op_type = local.operation(oai_deleted)
            if oai_deleted:
                # up to date only if the local record is deleted (or absent) as well
                if op_type is Operation.NOOP:
                    stream_entry.entry["record"] = None
                    return True
                return False
            if op_type is not Operation.UPDATE or local.draft is not None:
                # local record is not live (deleted, draft only, missing) or has a local
                # edit draft - write it (the draft gets overwritten and published)
                return False
            # try to read the existing record
            try:
                fetched_record = self.service.read(system_identity, oai_record.record_pid)
                stream_entry.entry["record"] = fetched_record.to_dict()
            except Exception:  # noqa BLE001 ok to catch here
                current_app.logger.warning(
                    "Failed to read existing record %s for OAI identifier %s. Re-writing the record.",
                    oai_record.record_pid,
                    oai_record.oai_identifier,
                )
                # proceed to re-write the record
            else:
                return True
        return False

    def _get_pid_value_from_record(self, record_data: dict) -> str | None:
        """Read the pid value from the transformed data via the pid provider, if it supports it."""
        pid_field = getattr(self.service.record_cls.pid, "field", None)
        pid_provider = getattr(pid_field, "_provider", None) if pid_field else None
        get_pid_value_from_record = getattr(pid_provider, "get_pid_value_from_record", None) if pid_provider else None
        pid_type = getattr(pid_field, "_pid_type", None) if pid_field else None
        if get_pid_value_from_record and pid_type:
            # try to get pid value from transformed data
            return get_pid_value_from_record(record_data)
        return None

    def _store_oai_record(  # noqa PLR0913,PLR0917
        self,
        oai_record: OAIHarvestedRecord | None,
        oai_identifier: str,
        oai_datestamp: dt.datetime,
        oai_deleted: bool,
        harvested_at: dt.datetime,
        original_data: dict,
        transformed_data: dict | None,
        stream_entry: StreamEntry,
        validation_errors: list[dict],
        exception_raised: Exception | None,
        written_data: dict | None,
    ) -> list[dict]:
        """Create or update the OAIHarvestedRecord instance, return the stored errors."""
        if not oai_record:
            oai_record = OAIHarvestedRecord(
                oai_identifier=oai_identifier,
                harvester_id=self._harvester_id,
            )
        oai_record.datestamp = oai_datestamp
        oai_record.deleted = oai_deleted
        # store errors and warnings
        # build the list before assigning - in-place append on a JSON column is not persisted
        errors = [x for e in (stream_entry.errors or []) for x in self._convert_stream_error(e)]
        errors.extend(validation_errors)
        if exception_raised is not None:
            errors.extend(self._convert_exception_to_errors(exception_raised))
        if stream_entry.exc:
            errors.extend(self._convert_exception_to_errors(stream_entry.exc))
        oai_record.errors = json.loads(json.dumps(errors, default=str))
        oai_record.harvested_at = harvested_at
        oai_record.has_errors = bool(oai_record.errors)
        # no warnings tracking for now
        oai_record.has_warnings = False

        # store the internal identifier
        # written_entry.entry is a RecordItem from the service, not plain record
        if written_data is not None:
            pid = dict_lookup(written_data, self._pid_field)
            oai_record.record_pid = pid
            oai_record.record_type = self._model

        # store the original and transformed data; values that are not json serializable
        # (dates, lazy strings, ...) would make the commit fail
        oai_record.original_data = original_data
        oai_record.transformed_data = json.loads(json.dumps(transformed_data or {}, default=str))

        db.session.add(oai_record)
        db.session.commit()

        # indexed right away - the bulk queue is processed only by a running celery beat,
        # so the administration would show the state of a previous harvest until then
        try:
            indexer = current_oai_record_service.indexer
            indexer.index(OAIRecordAggregate.from_model(oai_record))
        except Exception:
            current_app.logger.exception(
                "Could not index harvest record %s, queued for bulk indexing", oai_record.oai_identifier
            )
            current_oai_record_service.indexer.bulk_index([oai_record.oai_identifier])
        return oai_record.errors

    def _write(
        self,
        record_data: dict,
        local: LocalRecord,
        op_type: Operation,
        oai_identifier: str,
        uow: UnitOfWork,
    ) -> dict | None:
        """Write the record data using the service.

        Operations other than `NOOP` and `CREATE` always have a pid value - a local
        record without a pid is classified as `ABSENT` (see
        [LocalRecord.resolve][oarepo_oaipmh_harvester.writers.oai_service.LocalRecord.resolve]).
        """
        match op_type:
            case Operation.NOOP:
                return None
            case Operation.CREATE:
                return cast("dict", self.service.create(self._identity, record_data, uow=uow).to_dict())
            case Operation.UPDATE_DRAFT:
                draft_service = cast("DraftRecordService", self.service)
                return cast(
                    "dict",
                    draft_service.update_draft(
                        self._identity, cast("str", local.pid_value), record_data, uow=uow
                    ).to_dict(),
                )
            case Operation.UPDATE | Operation.RESTORE | Operation.REVIVE:
                self._bring_back(local, op_type, uow)
                return self._update(record_data, local, uow)
            case Operation.DELETE | Operation.DELETE_RECORD | Operation.DELETE_DRAFT:
                self._delete(local, op_type, oai_identifier, uow)
                return None
        raise ValueError(f"Unknown operation type: {op_type}")

    def _bring_back(self, local: LocalRecord, op_type: Operation, uow: UnitOfWork) -> None:
        """Make a locally deleted record live again before it is updated (no-op for `UPDATE`)."""
        pid_value = cast("str", local.pid_value)
        if op_type is Operation.RESTORE:
            rdm_service = cast("RDMRecordService", self.service)
            if getattr(local.record, "deletion_status", None) == RecordDeletionStatusEnum.MARKED:
                rdm_service.unmark_record_for_purge(self._identity, pid_value, uow=uow)
            rdm_service.restore_record(self._identity, pid_value, uow=uow)
        elif op_type is Operation.REVIVE:
            self._revive(local)

    def _delete(self, local: LocalRecord, op_type: Operation, oai_identifier: str, uow: UnitOfWork) -> None:
        """Delete the local record (`DELETE` / `DELETE_RECORD`) and its local draft."""
        pid_value = cast("str", local.pid_value)
        if op_type is Operation.DELETE:
            self.service.delete(self._identity, pid_value, uow=uow)
        elif op_type is Operation.DELETE_RECORD:
            cast("RDMRecordService", self.service).delete_record(
                self._identity,
                pid_value,
                {"note": f"Deleted at the OAI-PMH source ({oai_identifier})", "is_visible": True},
                uow=uow,
            )
        self._delete_local_draft(local, uow)

    def _delete_local_draft(self, local: LocalRecord, uow: UnitOfWork) -> None:
        """Delete a local edit draft of a record deleted at the source.

        `delete_record` (now or when the record was tombstoned locally) removes the draft
        from the index but keeps it in the db, while `delete_draft` fails if the draft is
        not indexed - so index it first (the ops run in order on commit).
        """
        if local.draft is not None:
            draft_service = cast("DraftRecordService", self.service)
            uow.register(RecordIndexOp(local.draft, indexer=draft_service.indexer))
            draft_service.delete_draft(self._identity, cast("str", local.pid_value), uow=uow)

    def _revive(self, local: LocalRecord) -> None:
        """Bring a hard deleted record back with its original pid.

        Hard delete sets `model.json` to `None`, but the previous revision still has
        the full json - reverting to it undeletes the record (`is_deleted` is
        `json is None`). The harvested data is applied by the following update.
        """
        record = cast("Record", local.record)
        pid = cast("PersistentIdentifier", local.deleted_pid)
        last_revision_id = next(
            (r.revision_id for r in reversed(record.revisions) if r.model is not None and r.model.json is not None),
            None,
        )
        if last_revision_id is None:
            raise ValueError(f"Hard deleted record {local.pid_value} has no revision to revive from.")
        pid.sync_status(PIDStatus.REGISTERED)
        record.revert(last_revision_id)

    def _update(self, record_data: dict, local: LocalRecord, uow: UnitOfWork) -> dict:
        """Update a live record; draft services do not allow direct update, so go through a draft."""
        pid_value = cast("str", local.pid_value)
        if not self._publish or self._direct_update:
            written = self.service.update(self._identity, pid_value, record_data, uow=uow).to_dict()
            # a local edit draft would otherwise survive (and overwrite the harvested data when published)
            self._delete_local_draft(local, uow)
            return cast("dict", written)
        draft_service = cast("DraftRecordService", self.service)
        # edit returns the existing draft if there is one - it gets overwritten by the harvested data
        draft_service.edit(self._identity, pid_value, uow=uow)
        return cast(
            "dict",
            draft_service.update_draft(self._identity, pid_value, record_data, uow=uow).to_dict(),
        )

    def write_many(self, stream_entries: list[StreamEntry], *args: Any, **kwargs: Any) -> list[StreamEntry]:
        """Write the input entries using a given service."""
        # For now, just call write() for each entry
        return [self.write(stream_entry, *args, **kwargs) for stream_entry in stream_entries]

    def _convert_exception_to_errors(self, exception: Exception) -> list[dict]:
        """Turn an exception into error dicts (type, message, location).

        Validation exceptions (marshmallow `ValidationError`, rdm
        `ValidationErrorWithMessageAsList`) carry structured `messages` - their `str()` is
        the repr of a list or dict. They give one error per message with the field path as
        the location, like the validation errors returned by the service. Other exceptions
        give a single error with the message and a trimmed traceback with the locals of each
        frame as the location.
        """
        messages = getattr(exception, "messages", None)
        if messages:
            if isinstance(exception, ma.ValidationError):
                messages = exception.normalized_messages()
            return [
                {"type": type(exception).__name__, "message": message, "location": field}
                for field, message in _flatten_messages(messages)
            ]
        return [self._convert_exception_to_error_dict(exception)]

    def _convert_exception_to_error_dict(self, exception: Exception) -> dict:
        """Turn an exception into an error dict with a trimmed traceback and its locals."""
        tb = exception.__traceback__
        formatted_tb = []
        while tb:
            frame = tb.tb_frame
            fn = (frame.f_code.co_filename or "").split("site-packages/")[-1]
            formatted_tb.append(f"{frame.f_code.co_name}@{fn}:{tb.tb_lineno}")
            for local_name, local_value in frame.f_locals.items():
                try:
                    local_value_str = repr(local_value)
                except Exception:  # noqa: BLE001 a broken repr must not prevent storing the error
                    local_value_str = f"<{type(local_value).__name__}, repr failed>"
                if len(local_value_str) > MAX_TRACEBACK_VARIABLE_LENGTH:
                    local_value_str = local_value_str[: MAX_TRACEBACK_VARIABLE_LENGTH - 3] + "..."
                formatted_tb.append(f"    {local_name} = {local_value_str}")
            tb = tb.tb_next

        return {
            "type": type(exception).__name__,
            # e.g. KeyError() or AssertionError without a message
            "message": str(exception) or f"{type(exception).__name__} (no message)",
            "location": "\n".join(formatted_tb),
        }

    def _convert_stream_error(self, error: Any) -> list[dict]:
        """Normalize a stream error (dict, exception or other) into error dicts."""
        if isinstance(error, dict):
            return [error]
        if isinstance(error, Exception):
            return self._convert_exception_to_errors(error)
        return [{"message": str(error) or "Unknown error (no message)"}]

    def _convert_to_oai_error(self, error: Any) -> Generator[dict]:
        """Convert a service error to an OAI-PMH error dict (type, location, message)."""
        if not isinstance(error, dict):
            yield {"message": str(error)}
        else:
            for field, message in _flatten_messages(error.get("messages") or [], error.get("field", "unknown")):
                yield {
                    "type": "validation_error",
                    "location": field,
                    "message": message,
                }