import concurrent.futures
from copy import copy, deepcopy
from datetime import datetime, timedelta
from enum import Enum
import os
from os import PathLike
from pathlib import Path
from typing import Union

from adcircpy import AdcircMesh
from nemspy import ModelingSystem

from coupledmodeldriver import Platform
from coupledmodeldriver.generate.adcirc.configure import (
    ADCIRCRunConfiguration,
    NEMSADCIRCRunConfiguration,
)
from coupledmodeldriver.generate.adcirc.script import (
    AdcircRunJob,
    AdcircSetupJob,
    AswipCommand,
    AdcircEnsembleRunScript,
    AdcircEnsembleCleanupScript,
)
from coupledmodeldriver.script import SlurmEmailType
from coupledmodeldriver.utilities import (
    create_symlink,
    get_logger,
    ProcessPoolExecutorStackTraced,
)

LOGGER = get_logger('cplmdldrv')


class RunPhase(Enum):
    """
    options for run phase of ADCIRC execuction
    """

    COLDSTART = 'coldstart'
    HOTSTART = 'hotstart'


def generate_adcirc_configuration(
    configuration_directory: PathLike,
    output_directory: PathLike = None,
    relative_paths: bool = False,
    overwrite: bool = False,
    parallel: bool = True,
):
    """
    generate an ADCIRC run configuration from JSON configuration files generated by ``initialize_adcirc``

    :param configuration_directory: path containing JSON configuration files
    :param output_directory: path to store generated configuration files
    :param relative_paths: whether to write relative paths in generated configuration files
    :param overwrite: whether to overwrite existing files
    :param parallel: generate configurations concurrently
    """

    start_time = datetime.now()

    if not isinstance(configuration_directory, Path):
        configuration_directory = Path(configuration_directory)

    if output_directory is None:
        output_directory = configuration_directory
    elif not isinstance(output_directory, Path):
        output_directory = Path(output_directory)

    output_directory = output_directory.resolve()
    if not output_directory.is_absolute():
        output_directory = output_directory.absolute()

    if not output_directory.exists():
        os.makedirs(output_directory, exist_ok=True)

    if configuration_directory.absolute().resolve() != Path.cwd():
        starting_directory = Path.cwd()
        os.chdir(configuration_directory)
        configuration_directory = Path.cwd()
    else:
        starting_directory = None

    use_nems = 'configure_nems.json' in [
        filename.name.lower() for filename in configuration_directory.iterdir()
    ]

    if use_nems:
        LOGGER.debug(f'generating NEMS configuration')
        base_configuration = NEMSADCIRCRunConfiguration.read_directory(configuration_directory)
    else:
        LOGGER.debug(f'generating ADCIRC-only configuration')
        base_configuration = ADCIRCRunConfiguration.read_directory(configuration_directory)

    base_configuration.move_paths(configuration_directory)

    platform = base_configuration['modeldriver']['platform']

    job_duration = base_configuration['slurm']['job_duration']
    slurm_account = base_configuration['slurm']['account']
    partition = base_configuration['slurm']['partition']
    email_type = base_configuration['slurm']['email_type']
    email_address = base_configuration['slurm']['email_address']

    original_fort13_filename = base_configuration['adcirc']['fort_13_path']
    original_fort14_filename = base_configuration['adcirc']['fort_14_path']
    adcirc_processors = base_configuration['adcirc']['processors']
    spinup_duration = base_configuration['adcirc']['tidal_spinup_duration']
    use_original_mesh = base_configuration['adcirc']['use_original_mesh']

    if original_fort14_filename is None or not original_fort14_filename.exists():
        raise FileNotFoundError(f'mesh XY not found at "{original_fort14_filename}"')

    local_fort13_filename = output_directory / 'fort.13'
    local_fort14_filename = output_directory / 'fort.14'
    local_fort15_filename = output_directory / 'fort.15'

    do_spinup = spinup_duration is not None

    run_phase = 'HOTSTART' if do_spinup else 'COLDSTART'

    if slurm_account is None:
        slurm_account = platform.value['slurm_account']

    ensemble_run_script_filename = output_directory / f'run_{platform.name.lower()}.sh'
    ensemble_cleanup_script_filename = output_directory / f'cleanup.sh'

    if use_original_mesh:
        LOGGER.info(
            f'using original mesh from "{os.path.relpath(original_fort14_filename.resolve(), Path.cwd())}"'
        )
        if original_fort13_filename.exists():
            create_symlink(original_fort13_filename, local_fort13_filename)
        create_symlink(original_fort14_filename, local_fort14_filename)
    elif overwrite or not local_fort14_filename.exists():
        LOGGER.info(
            f'rewriting original mesh to "{os.path.relpath(local_fort14_filename.resolve(), Path.cwd())}"'
        )
        adcircpy_driver = base_configuration.adcircpy_driver
        try:
            adcircpy_driver.write(
                output_directory,
                overwrite=overwrite,
                fort13='fort.13' if original_fort13_filename.exists() else None,
                fort14='fort.14',
                fort15='fort.15',
                fort22=None,
                coldstart=None,
                hotstart=None,
                driver=None,
            )
        except Exception as error:
            LOGGER.warning(error)

    if local_fort15_filename.exists():
        os.remove(local_fort15_filename)

    if local_fort13_filename.exists():
        base_configuration['adcirc']['fort_13_path'] = local_fort13_filename
    if local_fort14_filename.exists():
        base_configuration['adcirc']['fort_14_path'] = local_fort14_filename

    runs_directory = output_directory / 'runs'
    if not runs_directory.exists():
        runs_directory.mkdir(parents=True, exist_ok=True)

    perturbations = base_configuration.perturb()

    LOGGER.info(
        f'generating {len(perturbations)} run configuration(s) in "{os.path.relpath(runs_directory.resolve(), Path.cwd())}"'
    )

    if parallel:
        process_pool = ProcessPoolExecutorStackTraced()
        LOGGER.info(f'leveraging {os.cpu_count()} processor(s)')
    else:
        process_pool = None

    futures = []

    if do_spinup:
        spinup_directory = output_directory / 'spinup'

        spinup_configuration = copy(base_configuration)

        if not spinup_configuration.files_exist(spinup_directory):
            spinup_kwargs = {
                'directory': spinup_directory,
                'configuration': spinup_configuration,
                'duration': spinup_duration,
                'local_fort13_filename': local_fort13_filename,
                'local_fort14_filename': local_fort14_filename,
                'link_mesh': True,
                'relative_paths': relative_paths,
                'overwrite': overwrite,
                'platform': platform,
                'adcirc_processors': adcirc_processors,
                'slurm_account': slurm_account,
                'job_duration': job_duration,
                'partition': partition,
                'email_type': email_type,
                'email_address': email_address,
            }

            if parallel:
                futures.append(process_pool.submit(write_spinup_directory, **spinup_kwargs))
            else:
                spinup_directory = write_spinup_directory(**spinup_kwargs)
                LOGGER.info(f'wrote configuration to "{spinup_directory}"')
    else:
        spinup_directory = None

    for run_name, run_configuration in perturbations.items():
        run_directory = runs_directory / run_name

        if not run_configuration.files_exist(run_directory):
            run_kwargs = {
                'directory': run_directory,
                'name': run_name,
                'phase': run_phase,
                'configuration': run_configuration,
                'local_fort13_filename': local_fort13_filename,
                'local_fort14_filename': local_fort14_filename,
                'link_mesh': use_original_mesh,
                'relative_paths': relative_paths,
                'overwrite': overwrite,
                'platform': platform,
                'adcirc_processors': adcirc_processors,
                'slurm_account': slurm_account,
                'job_duration': job_duration,
                'partition': partition,
                'email_type': email_type,
                'email_address': email_address,
                'do_spinup': do_spinup,
                'spinup_directory': spinup_directory,
            }

            if parallel:
                # destroy stored copy of adcircpy mesh, because it cannot be pickled across processes
                run_configuration.adcircpy_mesh = None
                futures.append(process_pool.submit(write_run_directory, **run_kwargs))
            else:
                write_run_directory(**run_kwargs)
                LOGGER.info(f'wrote configuration to "{run_directory}"')

    if parallel:
        for completed_future in concurrent.futures.as_completed(futures):
            LOGGER.info(f'wrote configuration to "{completed_future.result()}"')

    cleanup_script = AdcircEnsembleCleanupScript()
    LOGGER.debug(
        f'writing cleanup script "{os.path.relpath(ensemble_cleanup_script_filename.resolve(), Path.cwd())}"'
    )
    cleanup_script.write(filename=ensemble_cleanup_script_filename, overwrite=overwrite)

    LOGGER.info(
        f'writing ensemble run script "{os.path.relpath(ensemble_run_script_filename.resolve(), Path.cwd())}"'
    )
    run_job_script = AdcircEnsembleRunScript(
        platform=platform,
        commands=[
            'echo deleting previous ADCIRC output',
            f'sh {ensemble_cleanup_script_filename.name}',
            'echo deleting previous ADCIRC logs',
            'rm spinup/*.log runs/*/*.log',
        ],
        run_spinup=do_spinup,
    )
    run_job_script.write(ensemble_run_script_filename, overwrite=overwrite)

    if starting_directory is not None:
        os.chdir(starting_directory)

    LOGGER.info(f'finished in {datetime.now() - start_time}')


def write_spinup_directory(
    directory: PathLike,
    configuration: Union[ADCIRCRunConfiguration, NEMSADCIRCRunConfiguration],
    duration: timedelta,
    local_fort14_filename: PathLike,
    local_fort13_filename: PathLike = None,
    link_mesh: bool = False,
    relative_paths: bool = False,
    overwrite: bool = False,
    platform: Platform = None,
    adcirc_processors: int = None,
    slurm_account: str = None,
    job_duration: timedelta = None,
    partition: str = None,
    email_type: SlurmEmailType = None,
    email_address: str = None,
) -> Path:
    """
    write an ADCIRC tidal spinup configuration to the given directory

    :param directory: path to output directory
    :param configuration: run configuration
    :param duration: tidal spinup duration
    :param local_fort14_filename: path to ``fort.14`` file
    :param local_fort13_filename: path to ``fort.13`` file
    :param link_mesh: whether to create a symbolic link from the original mesh to this configuration
    :param relative_paths: whether to write relative paths (instead of absolute)
    :param overwrite: whether to overwrite existing files
    :param platform: HPC platform to configure for
    :param adcirc_processors: number of processors to assign to ADCIRC execution
    :param slurm_account: Slurm account name
    :param job_duration: wall-clock duration of Slurm job
    :param partition: Slurm partition name
    :param email_type: Slurm email notification option
    :param email_address: Slurm email notification address
    """

    if not isinstance(directory, Path):
        directory = Path(directory)
    if local_fort13_filename is not None and not isinstance(local_fort13_filename, Path):
        local_fort13_filename = Path(local_fort13_filename)

    if not directory.exists():
        directory.mkdir(parents=True, exist_ok=True)

    if not overwrite and configuration.files_exist(directory):
        LOGGER.warning(f'skipping existing directory "{directory}"')
        return directory

    setup_job_name = 'ADCIRC_SETUP_SPINUP'
    job_name = 'ADCIRC_COLDSTART_SPINUP'

    adcircpy_driver = configuration.adcircpy_driver

    if relative_paths:
        configuration.relative_to(directory, inplace=True)

    if 'nems' in configuration:
        nems = configuration.nemspy_modeling_system
        nems = ModelingSystem(
            nems.start_time - duration,
            nems.start_time,
            nems.interval,
            ocn=deepcopy(nems['OCN']),
            **nems.attributes,
        )
        processors = nems.processors
        model_executable = configuration['nems']['executable_path']
    else:
        nems = None
        processors = adcirc_processors
        model_executable = configuration['adcirc']['adcirc_executable_path']

    adcprep_path = configuration['adcirc']['adcprep_executable_path']
    aswip_path = configuration['adcirc']['aswip_executable_path']
    source_filename = configuration['adcirc']['source_filename']

    model_executable = update_path_relative(model_executable, relative_paths, directory)
    adcprep_path = update_path_relative(adcprep_path, relative_paths, directory)
    aswip_path = update_path_relative(aswip_path, relative_paths, directory)
    source_filename = update_path_relative(source_filename, relative_paths, directory)

    setup_script_filename = directory / 'setup.job'
    job_script_filename = directory / 'adcirc.job'

    if configuration.use_aswip:
        aswip_command = AswipCommand(path=aswip_path, nws=configuration['besttrack']['nws'])
    else:
        aswip_command = None

    setup_script = AdcircSetupJob(
        platform=platform,
        adcirc_mesh_partitions=adcirc_processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_partition=partition,
        slurm_run_name=setup_job_name,
        adcprep_path=adcprep_path,
        aswip_command=aswip_command,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{setup_job_name}.err.log',
        slurm_log_filename=f'{setup_job_name}.out.log',
        source_filename=source_filename,
    )

    job_script = AdcircRunJob(
        platform=platform,
        slurm_tasks=processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_run_name=job_name,
        executable=model_executable,
        slurm_partition=partition,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{job_name}.err.log',
        slurm_log_filename=f'{job_name}.out.log',
        source_filename=source_filename,
    )

    setup_script.write(setup_script_filename, overwrite=overwrite)
    job_script.write(job_script_filename, overwrite=overwrite)

    if 'nems' in configuration:
        LOGGER.debug(f'setting spinup to {duration}')

        nems.write(
            directory, overwrite=overwrite, include_version=True,
        )
        LOGGER.info(
            f'writing NEMS+ADCIRC tidal spinup configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    else:
        LOGGER.debug(
            f'writing ADCIRC tidal spinup configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    adcircpy_driver.write(
        directory,
        overwrite=overwrite,
        fort13=None if link_mesh else 'fort.13',
        fort14=None if link_mesh else 'fort.14',
        fort22='fort.22' if 'besttrack' in configuration else None,
        coldstart='fort.15',
        hotstart=None,
        driver=None,
    )
    if link_mesh:
        if local_fort13_filename is not None and local_fort13_filename.exists():
            create_symlink(local_fort13_filename, directory / 'fort.13', relative=True)
        create_symlink(local_fort14_filename, directory / 'fort.14', relative=True)

    return directory


def write_run_directory(
    directory: PathLike,
    name: str,
    phase: str,
    configuration: Union[ADCIRCRunConfiguration, NEMSADCIRCRunConfiguration],
    local_fort14_filename: PathLike,
    local_fort13_filename: PathLike = None,
    link_mesh: bool = False,
    relative_paths: bool = False,
    overwrite: bool = False,
    platform: Platform = None,
    adcirc_processors: int = None,
    slurm_account: str = None,
    job_duration: timedelta = None,
    partition: str = None,
    email_type: SlurmEmailType = None,
    email_address: str = None,
    do_spinup: bool = False,
    spinup_directory: PathLike = None,
) -> Path:
    """
    write an ADCIRC tidal spinup configuration to the given directory

    :param directory: path to output directory
    :param name: run name
    :param phase: ADCIRC run phase
    :param configuration: run configuration
    :param local_fort14_filename: path to ``fort.14`` file
    :param local_fort13_filename: path to ``fort.13`` file
    :param link_mesh: whether to create a symbolic link from the original mesh to this configuration
    :param relative_paths: whether to write relative paths (instead of absolute)
    :param overwrite: whether to overwrite existing files
    :param platform: HPC platform to configure for
    :param adcirc_processors: number of processors to assign to ADCIRC execution
    :param slurm_account: Slurm account name
    :param job_duration: wall-clock duration of Slurm job
    :param partition: Slurm partition name
    :param email_type: Slurm email notification option
    :param email_address: Slurm email notification address
    :param do_spinup: whether to use spinup results to hotstart run
    :param spinup_directory: path to spinup directory
    """

    if not isinstance(directory, Path):
        directory = Path(directory)
    if spinup_directory is not None and not isinstance(spinup_directory, Path):
        spinup_directory = Path(spinup_directory)
    if not isinstance(local_fort14_filename, Path):
        local_fort14_filename = Path(local_fort14_filename)
    if local_fort13_filename is not None and not isinstance(local_fort13_filename, Path):
        local_fort13_filename = Path(local_fort13_filename)

    if not directory.exists():
        directory.mkdir(parents=True, exist_ok=True)
    LOGGER.debug(
        f'writing run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
    )

    if not overwrite and configuration.files_exist(directory):
        LOGGER.warning(f'skipping existing directory "{directory}"')
        return directory

    setup_job_name = f'ADCIRC_SETUP_{name}'
    job_name = f'ADCIRC_{phase}_{name}'

    adcircpy_driver = configuration.adcircpy_driver

    if relative_paths:
        configuration.relative_to(directory, inplace=True)

    if 'nems' in configuration:
        nems = configuration.nemspy_modeling_system
        processors = nems.processors
        model_executable = configuration['nems']['executable_path']
    else:
        nems = None
        processors = adcirc_processors
        model_executable = configuration['adcirc']['adcirc_executable_path']

    adcprep_path = configuration['adcirc']['adcprep_executable_path']
    aswip_path = configuration['adcirc']['aswip_executable_path']
    source_filename = configuration['adcirc']['source_filename']

    model_executable = update_path_relative(model_executable, relative_paths, directory)
    adcprep_path = update_path_relative(adcprep_path, relative_paths, directory)
    aswip_path = update_path_relative(aswip_path, relative_paths, directory)
    source_filename = update_path_relative(source_filename, relative_paths, directory)

    setup_script_filename = directory / 'setup.job'
    job_script_filename = directory / 'adcirc.job'

    if configuration.use_aswip:
        aswip_command = AswipCommand(path=aswip_path, nws=configuration['besttrack']['nws'])
    else:
        aswip_command = None

    setup_script = AdcircSetupJob(
        platform=platform,
        adcirc_mesh_partitions=adcirc_processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_partition=partition,
        slurm_run_name=setup_job_name,
        adcprep_path=adcprep_path,
        aswip_command=aswip_command,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{setup_job_name}.err.log',
        slurm_log_filename=f'{setup_job_name}.out.log',
        source_filename=source_filename,
    )

    job_script = AdcircRunJob(
        platform=platform,
        slurm_tasks=processors,
        slurm_account=slurm_account,
        slurm_duration=job_duration,
        slurm_run_name=job_name,
        executable=model_executable,
        slurm_partition=partition,
        slurm_email_type=email_type,
        slurm_email_address=email_address,
        slurm_error_filename=f'{job_name}.err.log',
        slurm_log_filename=f'{job_name}.out.log',
        source_filename=source_filename,
    )

    setup_script.write(setup_script_filename, overwrite=overwrite)
    job_script.write(job_script_filename, overwrite=overwrite)

    if 'nems' in configuration:
        nems.write(
            directory, overwrite=overwrite, include_version=True,
        )
        LOGGER.info(
            f'writing NEMS+ADCIRC run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    else:
        LOGGER.debug(
            f'writing ADCIRC run configuration to "{os.path.relpath(directory.resolve(), Path.cwd())}"'
        )
    adcircpy_driver.write(
        directory,
        overwrite=overwrite,
        fort13=None if link_mesh else 'fort.13',
        fort14=None if link_mesh else 'fort.14',
        fort22='fort.22' if 'besttrack' in configuration else None,
        coldstart=None,
        hotstart='fort.15',
        driver=None,
    )

    if not link_mesh and local_fort14_filename.exists():
        link_mesh = configuration.adcircpy_mesh == AdcircMesh.open(local_fort14_filename)

    if link_mesh:
        if local_fort13_filename is not None and local_fort13_filename.exists():
            create_symlink(local_fort13_filename, directory / 'fort.13', relative=True)
        create_symlink(local_fort14_filename, directory / 'fort.14', relative=True)

    if do_spinup:
        for hotstart_filename in ['fort.67.nc', 'fort.68.nc']:
            try:
                create_symlink(
                    spinup_directory / hotstart_filename,
                    directory / hotstart_filename,
                    relative=True,
                )
            except:
                LOGGER.warning(
                    f'unable to link `{hotstart_filename}` from coldstart to hotstart; '
                    'you must manually link or copy this file after coldstart completes'
                )

    return directory


def update_path_relative(
    path: PathLike, relative: bool = False, relative_directory: PathLike = None
) -> Path:
    """
    update the given path based on the given relative path

    :param path: input path
    :param relative: whether to use a relative path to the input path
    :param relative_directory: directory to use to update input directory
    :return: updated path
    """

    if path is not None:
        if not isinstance(path, Path):
            path = Path(path)
        if relative_directory is None:
            relative_directory = Path.cwd()
        if relative:
            if not isinstance(relative_directory, Path):
                relative_directory = Path(relative_directory)
            if path.is_absolute():
                path = Path(os.path.relpath(path, relative_directory))
        elif not path.is_absolute():
            path = (relative_directory / path).resolve().absolute()
    return path
