"""File I/O related firetasks for writing, deleting, transferring and archiving files."""

from __future__ import annotations

import os
import shutil
import time
import traceback
from os.path import abspath, expanduser, expandvars
from typing import Any

from monty.shutil import compress_dir, decompress_dir

from fireworks.core.firework import FiretaskBase

__author__ = "Anubhav Jain, David Waroquiers, Shyue Ping Ong"
__copyright__ = "Copyright 2013, The Materials Project"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "Jan 6, 2014"


class FileWriteTask(FiretaskBase):
    """A Firetask to write files.

    Required params:
        - files_to_write: ([{filename:(str), contents:(str)}]) List of dicts with filenames
            and contents

    Optional params:
        - dest: (str) Shared path for files
    """

    _fw_name = "FileWriteTask"
    required_params = ["files_to_write"]
    optional_params = ["dest"]

    def run_task(self, fw_spec: dict[str, Any]) -> None:
        pth = self.get("dest", os.getcwd())
        for file_record in self["files_to_write"]:
            with open(os.path.join(pth, file_record["filename"]), "w") as fptr:
                fptr.write(file_record["contents"])


class FileDeleteTask(FiretaskBase):
    """A Firetask to delete files.

    Required params:
        - files_to_delete: ([str]) Filenames to delete

    Optional params:
        - dest: (str) Shared path for files
        - ignore_errors (bool): Whether to ignore errors. Defaults to True.
    """

    _fw_name = "FileDeleteTask"
    required_params = ["files_to_delete"]
    optional_params = ["dest", "ignore_errors"]

    def run_task(self, fw_spec: dict[str, Any]) -> None:
        pth = self.get("dest", os.getcwd())
        ignore_errors = self.get("ignore_errors", True)
        for file_name in self["files_to_delete"]:
            try:
                os.remove(os.path.join(pth, file_name))
            except Exception as ex:
                if not ignore_errors:
                    raise OSError(str(ex))


class FileTransferTask(FiretaskBase):
    """A Firetask to Transfer files. Note that.

    Required params:
        - mode: (str) - move, mv, copy, cp, copy2, copytree, copyfile, rtransfer
        - files: ([str]) or ([(str, str)]) - list of source files, or dictionary containing
                'src' and 'dest' keys

    Optional params:
        - dest: (str) destination directory, if not specified within files parameter (else optional)
        - server: (str) server host for remote transfer
        - user: (str) user to authenticate with on remote server
        - key_filename: (str) optional SSH key location for remote transfer
        - max_retry: (int) number of times to retry failed transfers; defaults to `0` (no retries)
        - retry_delay: (int) number of seconds to wait between retries; defaults to `10`
        - ignore_errors (bool): Optional. Whether to ignore errors. Defaults to False.
    """

    _fw_name = "FileTransferTask"
    required_params = ["mode", "files"]
    optional_params = ["dest", "server", "user", "key_filename", "max_retry", "retry_delay", "ignore_errors"]

    fn_list = {
        "move": shutil.move,
        "mv": shutil.move,
        "copy": shutil.copy,
        "cp": shutil.copy,
        "copy2": shutil.copy2,
        "copytree": shutil.copytree,
        "copyfile": shutil.copyfile,
    }

    def run_task(self, fw_spec: dict[str, Any]) -> None:
        shell_interpret = self.get("shell_interpret", True)
        ignore_errors = self.get("ignore_errors", False)
        max_retry = self.get("max_retry", 0)
        retry_delay = self.get("retry_delay", 10)
        mode = self.get("mode", "move")

        if mode == "rtransfer":
            # remote transfers
            # Create SFTP connection
            import paramiko

            ssh = paramiko.SSHClient()
            ssh.load_host_keys(expanduser(os.path.join("~", ".ssh", "known_hosts")))
            ssh.connect(self["server"], username=self.get("user"), key_filename=self.get("key_filename"))
            sftp = ssh.open_sftp()

        for file_entry in self["files"]:
            try:
                if "src" in file_entry:
                    src = (
                        os.path.abspath(expanduser(expandvars(file_entry["src"])))
                        if shell_interpret
                        else file_entry["src"]
                    )
                else:
                    src = abspath(expanduser(expandvars(file_entry))) if shell_interpret else file_entry

                if mode == "rtransfer":
                    dest = self["dest"]
                    if os.path.isdir(src):
                        if not self._rexists(sftp, dest):
                            sftp.mkdir(dest)

                        for child_name in os.listdir(src):
                            if os.path.isfile(os.path.join(src, child_name)):
                                sftp.put(os.path.join(src, child_name), os.path.join(dest, child_name))
                    else:
                        if not self._rexists(sftp, dest):
                            sftp.mkdir(dest)

                        sftp.put(src, os.path.join(dest, os.path.basename(src)))

                else:
                    if "dest" in file_entry:
                        dest = (
                            abspath(expanduser(expandvars(file_entry["dest"])))
                            if shell_interpret
                            else file_entry["dest"]
                        )
                    else:
                        dest = abspath(expanduser(expandvars(self["dest"]))) if shell_interpret else self["dest"]
                    FileTransferTask.fn_list[mode](src, dest)

            except Exception:
                traceback.print_exc()
                if max_retry:
                    # we want to avoid hammering either the local or remote machine
                    time.sleep(retry_delay)
                    self["max_retry"] -= 1
                    self.run_task(fw_spec)

                elif not ignore_errors:
                    raise ValueError(
                        f"There was an error performing operation {mode} from {self['files']} to {self['dest']}"
                    )

        if mode == "rtransfer":
            sftp.close()
            ssh.close()

    @staticmethod
    def _rexists(sftp, path) -> bool:
        """os.path.exists for paramiko's SCP object."""
        try:
            sftp.stat(path)
        except OSError as e:
            if e[0] == 2:
                return False
            raise
        else:
            return True


class CompressDirTask(FiretaskBase):
    """Compress all files in a directory.

    Args:
        dest (str): Optional. Path to compress.
        compression (str): Optional. Can only be gz or bz2. Defaults to gz.
        ignore_errors (bool): Optional. Whether to ignore errors. Defaults to False.
    """

    _fw_name = "CompressDirTask"
    optional_params = ["compression", "dest", "ignore_errors"]

    def run_task(self, fw_spec) -> None:
        ignore_errors = self.get("ignore_errors", False)
        dest = self.get("dest", os.getcwd())
        compression = self.get("compression", "gz")
        try:
            compress_dir(dest, compression=compression)
        except Exception:
            if not ignore_errors:
                raise ValueError(f"There was an error performing {compression=} in {dest}.")


class DecompressDirTask(FiretaskBase):
    """Decompress all files in a directory. Autodetects gz, bz2 and z file
    extensions.

    Args:
        dest (str): Optional. Path to decompress.
        ignore_errors (bool): Optional. Whether to ignore errors. Defaults to False.
    """

    _fw_name = "DecompressDirTask"
    optional_params = ["dest", "ignore_errors"]

    def run_task(self, fw_spec) -> None:
        ignore_errors = self.get("ignore_errors", False)
        dest = self.get("dest", os.getcwd())
        try:
            decompress_dir(dest)
        except Exception:
            if not ignore_errors:
                raise ValueError(f"There was an error performing decompression in {dest}.")


class ArchiveDirTask(FiretaskBase):
    """Wrapper around shutil.make_archive to make tar archives.

    Args:
        base_name (str): Name of the file to create, including the path,
            minus any format-specific extension.
        format (str): Optional. one of "zip", "tar", "bztar" or "gztar". Defaults to gztar.
    """

    _fw_name = "ArchiveDirTask"
    required_params = ["base_name"]
    optional_params = ["format"]

    def run_task(self, fw_spec) -> None:
        shutil.make_archive(self["base_name"], format=self.get("format", "gztar"), root_dir=".")
