import logging
import pathlib
from concurrent.futures import Future
from typing import List, Optional, Union

from nuplan.common.utils.s3_utils import is_s3_path
from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario
from nuplan.planning.simulation.callback.abstract_callback import AbstractCallback
from nuplan.planning.simulation.history.simulation_history import SimulationHistory, SimulationHistorySample
from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner
from nuplan.planning.simulation.simulation_log import SimulationLog
from nuplan.planning.simulation.simulation_setup import SimulationSetup
from nuplan.planning.simulation.trajectory.abstract_trajectory import AbstractTrajectory
from nuplan.planning.utils.multithreading.worker_pool import Task, WorkerPool

logger = logging.getLogger(__name__)


def _save_log_to_file(
    file_name: pathlib.Path, scenario: AbstractScenario, planner: AbstractPlanner, history: SimulationHistory
) -> None:
    """
    Create SimulationLog and save it to disk.
    :param file_name: to write to.
    :param scenario: to store in the log.
    :param planner: to store in the log.
    :param history: to store in the log.
    """
    simulation_log = SimulationLog(file_path=file_name, scenario=scenario, planner=None, simulation_history=history)
    simulation_log.save_to_file()


class SimulationLogCallback(AbstractCallback):
    """
    Callback for simulation logging/object serialization to disk.
    """

    def __init__(
        self,
        output_directory: Union[str, pathlib.Path],
        simulation_log_dir: Union[str, pathlib.Path],
        serialization_type: str,
        worker_pool: Optional[WorkerPool] = None,
    ):
        """
        Construct simulation log callback.
        :param output_directory: where scenes should be serialized.
        :param simulation_log_dir: Folder where to save simulation logs.
        :param serialization_type: A way to serialize output, options: ["json", "pickle", "msgpack"].
        """
        available_formats = ["pickle", "msgpack"]
        if serialization_type not in available_formats:
            raise ValueError(
                "The simulation log callback will not store files anywhere!"
                f"Choose at least one format from {available_formats} instead of {serialization_type}!"
            )

        self._output_directory = pathlib.Path(output_directory) / simulation_log_dir
        self._serialization_type = serialization_type
        if serialization_type == "pickle":
            file_suffix = '.pkl.xz'
        elif serialization_type == "msgpack":
            file_suffix = '.msgpack.xz'
        else:
            raise ValueError(f"Unknown option: {serialization_type}")
        self._file_suffix = file_suffix

        self._pool = worker_pool
        self._futures: List[Future[None]] = []

    @property
    def futures(self) -> List[Future[None]]:
        """
        Returns a list of futures, eg. for the main process to block on.
        :return: any futures generated by running any part of the callback asynchronously.
        """
        return self._futures

    def on_initialization_start(self, setup: SimulationSetup, planner: AbstractPlanner) -> None:
        """
        Create directory at initialization
        :param setup: simulation setup
        :param planner: planner before initialization
        """
        scenario_directory = self._get_scenario_folder(planner.name(), setup.scenario)

        if not is_s3_path(scenario_directory):
            scenario_directory.mkdir(exist_ok=True, parents=True)

    def on_initialization_end(self, setup: SimulationSetup, planner: AbstractPlanner) -> None:
        """Inherited, see superclass."""
        pass

    def on_step_start(self, setup: SimulationSetup, planner: AbstractPlanner) -> None:
        """Inherited, see superclass."""
        pass

    def on_step_end(self, setup: SimulationSetup, planner: AbstractPlanner, sample: SimulationHistorySample) -> None:
        """Inherited, see superclass."""
        pass

    def on_planner_start(self, setup: SimulationSetup, planner: AbstractPlanner) -> None:
        """Inherited, see superclass."""
        pass

    def on_planner_end(self, setup: SimulationSetup, planner: AbstractPlanner, trajectory: AbstractTrajectory) -> None:
        """Inherited, see superclass."""
        pass

    def on_simulation_start(self, setup: SimulationSetup) -> None:
        """Inherited, see superclass."""
        pass

    def on_simulation_end(self, setup: SimulationSetup, planner: AbstractPlanner, history: SimulationHistory) -> None:
        """
        On reached_end validate that all steps were correctly serialized.
        :param setup: simulation setup.
        :param planner: planner when simulation ends.
        :param history: resulting from simulation.
        """
        number_of_scenes = len(history)
        if number_of_scenes == 0:
            raise RuntimeError("Number of scenes has to be greater than 0")

        # Create directory
        scenario_directory = self._get_scenario_folder(planner.name(), setup.scenario)

        scenario = setup.scenario
        file_name = scenario_directory / (scenario.scenario_name + self._file_suffix)
        if self._pool is not None:
            self._futures = []
            self._futures.append(
                self._pool.submit(
                    Task(_save_log_to_file, num_cpus=1, num_gpus=0), file_name, scenario, planner, history
                )
            )
        else:
            _save_log_to_file(file_name, scenario, planner, history)

    def _get_scenario_folder(self, planner_name: str, scenario: AbstractScenario) -> pathlib.Path:
        """
        Compute scenario folder directory where all files will be stored.
        :param planner_name: planner name.
        :param scenario: for which to compute directory name.
        :return directory path.
        """
        return self._output_directory / planner_name / scenario.scenario_type / scenario.log_name / scenario.scenario_name  # type: ignore
