# -*- coding: utf-8 -*-
#
# Copyright (C) 2022-2023 CERN.
#
# Invenio-RDM-Migrator is free software; you can redistribute it and/or modify
# it under the terms of the MIT License; see LICENSE file for more details.

"""Invenio RDM migration files load module."""

from ...load.postgresql.bulk import PostgreSQLCopyLoad
from ...load.postgresql.bulk.generators import ExistingDataTableGenerator
from ..models.files import FilesBucket, FilesInstance, FilesObjectVersion
from .table_generator import FilesTableGenerator


class FilesCopyLoad(PostgreSQLCopyLoad):
    """PostgreSQL files COPY load."""

    def __init__(self, **kwargs):
        """Constructor."""
        super().__init__(table_generators=[FilesTableGenerator()], **kwargs)


class ExistingFilesLoad(PostgreSQLCopyLoad):
    """Existing files class for data loading."""

    def __init__(self, db_uri, data_dir, **kwargs):
        """Constructor."""
        super().__init__(
            db_uri=db_uri,
            table_generators=[
                ExistingDataTableGenerator(
                    tables=[FilesInstance, FilesBucket, FilesObjectVersion], pks=[]
                )
            ],
            data_dir=data_dir,
            existing_data=True,
        )
