from nmdc_schema.migrators.migrator_base import MigratorBase

class Migrator(MigratorBase):
    r"""Migrates a database between two schemas."""

    _from_version = "11.8.0"
    _to_version = "11.9.0.part_1"

    wfe_mappings = {}
    data_objects_mapping = {}

    def upgrade(self) -> None:
        r"""Migrates the database from conforming to the original schema, to conforming to the new schema.
        This migrator addresses issue https://github.com/microbiomedata/nmdc-schema/issues/1727.
        This will migrate fields from all `WorkflowExecution` records with type `nmdc:MetagenomeSequencing` to related `DataGeneration` records. Deletion of the `MetagenomeSequencing` records will occur outside of this migrator via the `/queries:run` API endpoint.
        
        The following fields will be migrated from the WorkflowExecution records to the DataGeneration records:
        - has_output
        
        Additionally, it will properly link the DataObjects to the associated DataGeneration records through the was_generated_by field.

        Important links to note:
        - The `was_informed_by` field in the WorkFlowExecution records will be used to link them to the DataGeneration records that were created to replace the deprecated WorkflowExecution records.
        - The `has_output` field in the WorkFlowExecution records will be used to identify the DataObjects that were generated by the records of interest, and will be added to the DataGeneration records.
        - the 'was_generated_by' field in the DataObjects will be updated to point to the corresponding DataGeneration records.
        """
        self.adapter.do_for_each_document("workflow_execution_set", self.store_we_ms_fields)
        self.adapter.process_each_document("data_object_set", [self.link_do_to_dg])
        self.adapter.process_each_document("data_generation_set", [self.migrate_fields_to_dg])


    def store_we_ms_fields(self, workflow_execution_record: dict) -> None:
        r"""
        For each workflow execution record that is of type MetagenomeSequencing, gather the following fields to fill the corresponding DataGeneration record:
        - has_output (used to fill the has_output field in the DataGeneration record, and used to link the DataObjects to the DataGeneration record)

        >>> m = Migrator()
        >>> m.store_we_ms_fields({"id": "nmdc:wfmsa-11-kd2tsp05.1", "type": "nmdc:Anything", "has_output": ["nmdc:dobj-11-nrjyjm34"], "was_informed_by": "nmdc:omprc-11-8yy07g22"})
        >>> m.wfe_mappings
        {}
        >>> m.data_objects_mapping
        {}
        >>> m.store_we_ms_fields({"id": "nmdc:wfmsa-11-kd2tsp04.1", "type": "nmdc:MetagenomeSequencing", "has_output": ["nmdc:dobj-11-nrjyjm33"], "was_informed_by": "nmdc:omprc-11-8yy07g21"})
        >>> m.wfe_mappings
        {'nmdc:omprc-11-8yy07g21': {'ms_id': 'nmdc:wfmsa-11-kd2tsp04.1', 'has_output': ['nmdc:dobj-11-nrjyjm33']}}
        >>> m.data_objects_mapping
        {'nmdc:dobj-11-nrjyjm33': 'nmdc:omprc-11-8yy07g21'}
        """
        if workflow_execution_record.get("type") == "nmdc:MetagenomeSequencing":
            ms_id = workflow_execution_record.get("id")
            dg_id = workflow_execution_record.get("was_informed_by")
            data_object_ids = workflow_execution_record.get("has_output", [])
            dg_record = {
                "ms_id": ms_id,
                "has_output": data_object_ids,
            }
            assert len(data_object_ids) == 1, "This migrator was implemented under the assumption that the `WorkflowExecution`'s `has_output` list contains exactly one item. A `has_output` list greater than 1 was found. Migrating this database will require modifying the migrator."
            self.data_objects_mapping[data_object_ids[0]] = workflow_execution_record.get("was_informed_by")
            # Add the dict to the wfe_mappings
            self.wfe_mappings[dg_id] = dg_record
            

    def link_do_to_dg(self, data_object_record: dict) -> dict:
        r"""
        For each data object, link it to the corresponding DataGeneration record through the was_generated_by field using the data_objects_mapping.
        
        >>> m = Migrator()
        >>> m.data_objects_mapping = {'nmdc:dobj-11-nrjyjm33': 'nmdc:omprc-11-8yy07g21'}
        >>> m.link_do_to_dg({"id": "nmdc:dobj-11-nrjyjm33", "was_generated_by": None})
        {'id': 'nmdc:dobj-11-nrjyjm33', 'was_generated_by': 'nmdc:omprc-11-8yy07g21'}
        >>> m.link_do_to_dg({"id": "nmdc:dobj-11-nrjyjm34", "was_generated_by": "nmdc:some_other_record"})
        {'id': 'nmdc:dobj-11-nrjyjm34', 'was_generated_by': 'nmdc:some_other_record'}
        """
        #create a list of the data object ids that we will be adding to
        data_object_ids = self.data_objects_mapping.keys()
        # get the id of the current data object
        do_id = data_object_record.get("id")
        # Check if the data object is an object of interest
        if do_id in data_object_ids:
            # Check if the data object already has a was_generated_by field that has values.
            if not data_object_record.get("was_generated_by"):
                data_object_record["was_generated_by"] = self.data_objects_mapping.get(do_id)
            else:
                # If it already has a was_generated_by field, we do not want to overwrite it
                self.logger.warning(f"Data object {do_id} already has a was_generated_by field, skipping update.")
        return data_object_record

    def migrate_fields_to_dg(self, data_generation_record: dict) -> dict:
        r"""
        Migrate the fields from the workflow execution record to the DataGeneration record.
        
        >>> m = Migrator()
        >>> m.wfe_mappings = {'nmdc:omprc-11-8yy07g21': {'ms_id': 'nmdc:wfmsa-11-kd2tsp04.1', 'has_output': ['nmdc:dobj-11-nrjyjm33']}}
        >>> m.migrate_fields_to_dg({"id": "nmdc:omprc-11-8yy07g21", "has_output": []})
        {'id': 'nmdc:omprc-11-8yy07g21', 'has_output': ['nmdc:dobj-11-nrjyjm33']}
        >>> m.migrate_fields_to_dg({"id": "nmdc:omprc-11-8yy07g22", "has_output": []})
        {'id': 'nmdc:omprc-11-8yy07g22', 'has_output': []}
        """
        # Get the target ids from the wfe_mappings
        target_ids = list(self.wfe_mappings.keys())
        dg_record_id = data_generation_record.get("id")

        # Check if the data generation record's id is one of the target ids
        if dg_record_id in target_ids:
            # get the corresponding mapping from wfe_mappings
            wfe_map = self.wfe_mappings[dg_record_id]
            dg_has_output = data_generation_record.pop("has_output", [])
            wfe_map_has_output = wfe_map.get("has_output", [])
            new_has_output = list(set(dg_has_output + wfe_map_has_output))
            # Migrate the has_output field
            data_generation_record["has_output"] = new_has_output
        return data_generation_record