import logging
from pathlib import Path
from typing import List

import argschema
from allensdk.brain_observatory.behavior.behavior_session import (
    BehaviorSession,
)
from allensdk.brain_observatory.behavior.write_nwb.behavior.schemas import (
    BehaviorInputSchema,
    OutputSchema,
)
from allensdk.brain_observatory.nwb.nwb_utils import NWBWriter


class WriteBehaviorNWB(argschema.ArgSchemaParser):
    default_schema = BehaviorInputSchema
    default_output_schema = OutputSchema

    def run(self):
        self.logger.name = type(self).__name__
        bs_id = self.args["behavior_session_id"]
        bs_id_dir = Path(self.args["output_dir_path"]) / f"{bs_id}"
        bs_id_dir.mkdir(exist_ok=True)
        output_file = self.write_behavior_nwb(
            behavior_session_metadata=self.args["behavior_session_metadata"],
            nwb_filepath=bs_id_dir / f"behavior_session_{bs_id}.nwb",
            skip_metadata=self.args["skip_metadata_key"],
            skip_stim=self.args["skip_stimulus_file_key"],
        )
        logging.info("File successfully created")

        output_dict = {
            "output_path": output_file,
            "input_parameters": self.args,
        }

        self.output(output_dict)

    def write_behavior_nwb(
        self,
        behavior_session_metadata: dict,
        nwb_filepath: Path,
        skip_metadata: List[str],
        skip_stim: List[str],
    ) -> str:
        """Load and write a BehaviorSession as NWB.

        Check data object against associated data from the behavior_session
        metadata table. Additionally, round trip the NWB file to confirm that
        it was properly created for the session.

        Parameters
        ----------
        behavior_session_metadata : dict
            Dictionary of keys that overlap between the metadata in the
            BehaviorSession object and the associated metadata table.
        nwb_filepath : pathlib.Path
            Base filename path to write the NWB to.
        skip_metadata : list of str
            List of metadata keys to skip when comparing data.
        skip_stim : list of str
            List of stimulus file keys to skip when comparing data.

        Returns
        -------
        output_path : str
            String path of where the NWB was written to.

        Note
        ----
        Upon encountering an error, either from write or data comparison, the
        code changes the file name suffix to error and does not return the
        path of the file.
        """

        nwb_writer = NWBWriter(
            nwb_filepath=str(nwb_filepath),
            session_data=behavior_session_metadata,
            serializer=BehaviorSession,
        )
        nwb_writer.write_nwb(skip_metadata=skip_metadata, skip_stim=skip_stim)

        return str(nwb_filepath)


if __name__ == "__main__":
    write_nwb = WriteBehaviorNWB()
    write_nwb.run()
