import logging
import re
from datetime import datetime, timedelta
from pathlib import Path

import f90nml

logger = logging.getLogger("prepare_chunk")
logger.setLevel(logging.INFO)

# Get some autosubmit variables
WORKDIR = "%HPCROOTDIR%"
STARTDATE = "%SDATE%"
MEMBER = "%MEMBER%"
CHUNK = "%CHUNK%"

# Example of date format "2018-06-01T00:00:00Z"
date_format = "%simulation.date_format%"

START_YEAR = "%Chunk_START_YEAR%"
START_MONTH = "%Chunk_START_MONTH%"
START_DAY = "%Chunk_START_DAY%"
START_HOUR = "%Chunk_START_HOUR%"

END_YEAR = "%Chunk_END_YEAR%"
END_MONTH = "%Chunk_END_MONTH%"
END_DAY = "%Chunk_END_DAY%"
END_HOUR = "%Chunk_END_HOUR%"

Chunk_START_DATE = datetime(year=int(START_YEAR), month=int(START_MONTH), day=int(START_DAY), hour=int(START_HOUR))
Chunk_END_DATE = datetime(year=int(END_YEAR), month=int(END_MONTH), day=int(END_DAY), hour=int(END_HOUR))

# Compute difference in seconds
checkpoint_time = int((Chunk_END_DATE - Chunk_START_DATE).total_seconds())

# TODO: Is that really necessary?
# Add 10 minutes to allow the model to write the restarts
Chunk_END_DATE = Chunk_END_DATE + timedelta(minutes=10)
# Get run directory
RUNDIR = Path(f"{WORKDIR}/{STARTDATE}/{MEMBER}")

# TODO: This is a bit ugly
# Read first-guess and analysis filenames from files:
first_guess_filename = (RUNDIR / "fg_file.txt").read_text().strip()
analysis_filename = (RUNDIR / "an_file.txt").read_text().strip()

# Get some variable replacements from the proj.yml file through autosubmit
variable_replacements = {
    "dynamics_grid_filename": "%simulation.dynamics_grid_filename%",
    "radiation_grid_filename": "%simulation.radiation_grid_filename%",
    "external_parameters_filename": "%simulation.external_parameters_filename%",
    "first_guess_filename": first_guess_filename,
    "analysis_filename": analysis_filename,
    "Chunk_START_DATE": Chunk_START_DATE.strftime(date_format),
    "Chunk_END_DATE": Chunk_END_DATE.strftime(date_format),
    "is_restart": False if "%CHUNK%" == "1" else True,
    "checkpoint_time": checkpoint_time,
}


def adapt_namelist(input_namelist: str, output_namelist: str):
    input_namelist = Path(input_namelist)
    output_namelist = Path(output_namelist)

    namelist = f90nml.read(input_namelist.as_posix())
    group_keys = [gk for gk in namelist]

    for group in group_keys:
        variable_keys = [vk for vk in namelist[group]]
        for variable in variable_keys:
            value = namelist[group][variable]
            m = re.match(r"%(.*)%", str(value))
            if m:
                key = m.group(1)

                if key not in variable_replacements:
                    raise AssertionError(f"The namelist {input_namelist.as_posix()!r} contains the variable {key!r} "
                                         f"which is not in the list of provided replacements:\n"
                                         f"{[v for v in variable_replacements]}")
                logger.info(f"Replacing {group}>{variable}:{key} with {variable_replacements[key]!r}")
                namelist[group][variable] = variable_replacements[key]

    f90nml.write(nml=namelist, nml_path=output_namelist.as_posix(), force=True)


if __name__ == '__main__':
    atmosphere_namelist_path = "%simulation.namelist_paths.atmosphere%"
    master_namelist_path = "%simulation.namelist_paths.master%"

    # Adapt atmosphere namelist
    adapt_namelist(input_namelist=atmosphere_namelist_path,
                   output_namelist=(RUNDIR / "icon_atmosphere.namelist").as_posix())
    # Adapt master namelist
    adapt_namelist(input_namelist=master_namelist_path,
                   output_namelist=(RUNDIR / "icon_master.namelist").as_posix())