from pathlib import Path

__all__ = [
    "init_input_node",
    "save_to_caps",
    "write_tsv_files",
]


def init_input_node(
    t1w: Path, recon_all_args: str, output_dir: Path
) -> tuple[str, Path, str, Path]:
    """Initialize the pipeline.

    This function will:
        - Extract <image_id> (e.g. sub-CLNC01_ses-M000) T1w filename;
        - Check FOV of T1w;
        - Create SUBJECTS_DIR for recon-all (otherwise, the command won't run);
        - Print begin execution message.
    """
    from clinica.utils.filemanip import get_subject_id
    from clinica.utils.ux import print_begin_image

    image_id = get_subject_id(t1w)
    flags = _check_flags(t1w, recon_all_args)
    subjects_dir = output_dir / image_id
    subjects_dir.mkdir(parents=True, exist_ok=True)
    print_begin_image(image_id, ["ReconAllArgs"], [flags])

    return image_id, t1w, flags, subjects_dir


def _check_flags(t1w_image: Path, recon_all_args: str) -> str:
    """Check `recon_all_args` flags for `in_t1w` image.

    Currently, this function only adds '-cw256' if the FOV of `in_t1w` is greater than 256.

    Parameters
    ----------
    t1w_image : str or Path
        The path to the T1W image.

    recon_all_args : str
        ??????.

    Returns
    -------
    str :
        ????.
    """
    import nibabel as nib

    image = nib.load(t1w_image)
    if any(
        [
            v * t > 256
            for v, t in zip(image.header.get_zooms(), image.header.get_data_shape())
        ]
    ):
        return " ".join([recon_all_args, "-cw256"])
    return recon_all_args


def write_tsv_files(subjects_dir: Path, image_id: str) -> str:
    """
    Generate statistics TSV files in `subjects_dir`/regional_measures folder for `image_id`.

    Notes
    -----
    We do not need to check the line "finished without error" in scripts/recon-all.log.
    If an error occurs, it will be detected by Nipype and the next nodes (including
    write_tsv_files will not be called).
    """
    from clinica.pipelines.anatomical.freesurfer.utils import generate_regional_measures
    from clinica.utils.stream import cprint

    if (subjects_dir / image_id / "mri" / "aparc+aseg.mgz").is_file():
        generate_regional_measures(subjects_dir, image_id)
    else:
        cprint(
            msg=(
                f"{image_id.replace('_', ' | ')} does not contain "
                f"mri/aseg+aparc.mgz file. Creation of regional_measures/ folder will be skipped."
            ),
            lvl="warning",
        )
    return image_id


def save_to_caps(
    source_dir: Path, image_id: str, caps_dir: Path, overwrite_caps: bool = False
) -> str:
    """Save `source_dir`/`image_id`/ to CAPS folder.

    This function copies FreeSurfer segmentation and regional_measures folder of `source_dir`/`image_id`/ to
    `caps_dir`/subjects/<participant_id>/<session_id>/t1_freesurfer_cross_sectional/
    where `image_id` = <participant_id>_<session_id>.

    The `source_dir`/`image_id`/ folder should contain the following elements:
        - fsaverage, lh.EC_average and rh.EC_average symbolic links automatically generated by recon-all
        - `image_id`/ folder containing the FreeSurfer segmentation
        - regional_measures/ folder containing TSV files

    Notes
    -----
    We do not need to check the line "finished without error" in scripts/recon-all.log.
    If an error occurs, it will be detected by Nipype and the next nodes (i.e.  save_to_caps will not be called).
    """
    import shutil

    from clinica.utils.stream import cprint
    from clinica.utils.ux import print_end_image

    participant_id = image_id.split("_")[0]
    session_id = image_id.split("_")[1]
    destination_dir = (
        caps_dir.expanduser()
        / "subjects"
        / participant_id
        / session_id
        / "t1"
        / "freesurfer_cross_sectional"
    )
    source_file = (
        source_dir.expanduser() / image_id / image_id / "mri" / "aparc+aseg.mgz"
    )
    destination_file = destination_dir / image_id / "mri" / "aparc+aseg.mgz"
    if source_file.is_file():
        if destination_file.is_file():
            if overwrite_caps:
                shutil.rmtree(destination_dir)
                shutil.copytree(
                    src=source_dir / image_id / image_id,
                    dst=destination_dir / image_id,
                    symlinks=True,
                )
                shutil.copytree(
                    src=source_dir / image_id / "regional_measures",
                    dst=destination_dir / "regional_measures",
                    symlinks=True,
                )
        else:
            shutil.copytree(
                src=source_dir / image_id / image_id,
                dst=destination_dir / image_id,
                symlinks=True,
            )
            shutil.copytree(
                src=source_dir / image_id / "regional_measures",
                dst=destination_dir / "regional_measures",
                symlinks=True,
            )
        print_end_image(image_id)
    else:
        cprint(
            msg=(
                f"{image_id.replace('_', ' | ')} does not "
                "contain mri/aseg+aparc.mgz file. Copy will be skipped."
            ),
            lvl="warning",
        )
    return image_id
