"""Module containing command line scripts for developers."""

import click


@click.group(context_settings=dict(help_option_names=["-h", "--help"]))
def dev():
    """Tools for atomate2 developers."""


@dev.command(context_settings=dict(help_option_names=["-h", "--help"]))
@click.argument("test_dir")
def vasp_test_data(test_dir):
    """Generate test data for VASP unit tests.

    This script expects there is an outputs.json file and job folders in the current
    directory. Please refer to the atomate2 documentation on writing unit tests for more
    information.
    """
    import sys
    import warnings
    from pathlib import Path
    from pprint import pformat

    from monty.serialization import loadfn

    from atomate2.common.files import copy_files, delete_files, gunzip_files
    from atomate2.utils.path import strip_hostname
    from atomate2.vasp.schemas.task import TaskDocument

    warnings.filterwarnings("ignore", module="pymatgen")

    test_dir = Path(test_dir)

    if test_dir.exists():
        click.echo("test_data folder already exists, refusing to overwrite it")
        sys.exit()

    test_dir.mkdir()

    outputs = loadfn("outputs.json")

    task_labels = [
        o["output"].task_label for o in outputs if isinstance(o, TaskDocument)
    ]

    if len(task_labels) != len(set(task_labels)):
        raise ValueError("Not all jobs have unique names")

    original_mapping = {}
    mapping = {}
    for output in outputs:
        if not isinstance(output["output"], TaskDocument):
            # this is not a VASP job
            continue

        job_name = output["output"].task_label
        orig_job_dir = strip_hostname(output["output"].dir_name)
        folder_name = output["output"].task_label.replace("/", "_").replace(" ", "_")

        if len(task_labels) == 1:
            # only testing a single job
            job_dir = test_dir
        else:
            job_dir = test_dir / folder_name
            job_dir.mkdir()

        mapping[job_name] = str(job_dir)
        original_mapping[str(job_dir)] = orig_job_dir

        # create input folder and copy across files
        input_dir = job_dir / "inputs"
        input_dir.mkdir()

        copy_files(
            orig_job_dir,
            input_dir,
            include_files=[
                "INCAR",
                "INCAR.gz",
                "KPOINTS",
                "KPOINTS.gz",
                "POTCAR",
                "POTCAR.gz",
                "POSCAR",
                "POSCAR.gz",
            ],
            allow_missing=True,
        )
        gunzip_files(input_dir)
        _potcar_to_potcar_spec(input_dir / "POTCAR", input_dir / "POTCAR.spec")
        delete_files(input_dir, include_files=["POTCAR"])

        # create output folder and copy across files
        output_dir = job_dir / "outputs"
        output_dir.mkdir()
        copy_files(
            orig_job_dir,
            output_dir,
            include_files=[
                "POSCAR*",
                "CONTCAR*",
                "KPOINTS*",
                "INCAR*",
                "vasprun*",
                "OUTCAR*",
                "*.json*",
            ],
            allow_missing=True,
        )
        copy_files(input_dir, output_dir, include_files=["POTCAR.spec"])

    mapping_str = pformat(mapping).replace("\n", "\n    ")
    original_mapping_str = "\n".join(
        [f"  {v}  ->  {k}" for k, v in original_mapping.items()]
    )

    run_vasp_kwargs = {k: {"incar_settings": ["NSW", "ISMEAR"]} for k in mapping.keys()}
    run_vasp_kwargs_str = pformat(run_vasp_kwargs).replace("\n", "\n    ")

    test_function_str = f"""Test files generated in test_data.

Please ensure that all other necessary files are included in the test_data, such as
CHGCAR, LOCPOT, etc. However, only include additional files if they are absolutely
necessary (for example to calculate defect corrections) otherwise they will increase
the size of the atomate2 repository.

A mapping from the original job folders to the formatted folders is:
{original_mapping_str}

An example test using the generated test data is provided below. Be sure to update
this test with the particulars for your system. For more examples, see the existing
tests in atomate2/tests/vasp/jobs.

def test_my_flow(mock_vasp, clean_dir, si_structure):
    from jobflow import run_locally

    # mapping from job name to directory containing test files
    ref_paths = {mapping_str}

    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
    fake_run_vasp_kwargs = {run_vasp_kwargs_str}

    # automatically use fake VASP and write POTCAR.spec during the test
    mock_vasp(ref_paths, fake_run_vasp_kwargs)

    # !!! Generate job
    job = MyMaker().make(si_structure)

    # run the flow or job and ensure that it finished running successfully
    responses = run_locally(job, create_folders=True, ensure_success=True)

    # !!! validation on the outputs
    output1 = responses[job.uuid][1].output
    assert isinstance(output1, TaskDocument)
    assert output1.output.energy == pytest.approx(-10.85037078)
    """

    print(test_function_str)


def _potcar_to_potcar_spec(potcar_filename, output_filename):
    """Convert a POTCAR file to a POTCAR.spec file."""
    from pymatgen.io.vasp import Potcar

    potcar = Potcar.from_file(potcar_filename)
    output_filename.write_text("\n".join(potcar.symbols))
