# (c) 2015-2023 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
"""Local filesystem manifest and execution backend."""

import datetime
import json
import logging
import os
import shutil
import stat
from glob import glob
from typing import Any, Dict, List, Optional
from types import ModuleType
from playmolecule._appfiles import _File
from playmolecule._config import _get_config
from playmolecule._validation import _VALIDATORS, _validate_argument_type
from playmolecule._public_api import JobStatus

logger = logging.getLogger(__name__)

JOB_TIMEOUT = 60  # Timeout in seconds for heartbeat


curr_dir = os.path.dirname(os.path.abspath(__file__))


class _LocalManifestBackend:
    """Backend for loading app manifests from local filesystem."""

    def __init__(self, root: str):
        """Initialize the local manifest backend.

        Parameters
        ----------
        root : str
            Root directory containing apps
        """
        self.root = root

    def get_apps(self) -> Dict[str, Dict[str, dict]]:
        """Get all available apps from local filesystem.

        Returns
        -------
        dict
            Nested dict: {appname: {version: {"manifest": ..., "appdir": ..., "run.sh": ...}}}
        """
        from natsort import natsorted

        apps = {}

        for app_d in natsorted(glob(os.path.join(self.root, "apps", "*", ""))):
            appname = os.path.basename(os.path.abspath(app_d))
            if not self._check_folder_validity(app_d):
                continue

            versions = glob(os.path.join(app_d, "*"))
            versions.sort(key=lambda s: natsorted(os.path.basename(s)))

            app_versions = {}
            for version_folder in versions:
                vname = os.path.basename(version_folder)
                if not self._check_folder_validity(version_folder):
                    continue

                run_sh = os.path.join(version_folder, "run.sh")
                if not os.path.exists(run_sh):
                    logger.error(f"Run script not found in {version_folder}")
                    continue
                with open(run_sh, "r") as f:
                    run_sh = f.read()

                jsonf = glob(os.path.join(version_folder, "*.json"))
                if len(jsonf) > 1:
                    logger.error(f"Multiple json files found in {version_folder}")
                    continue
                if len(jsonf) == 1 and os.stat(jsonf[0]).st_size != 0:
                    try:
                        with open(jsonf[0], "r") as f:
                            manifest = json.load(f)
                        app_versions[vname] = {
                            "manifest": manifest,
                            "appdir": version_folder,
                            "run.sh": run_sh,
                            "files": self._get_app_files(manifest, version_folder),
                        }
                    except Exception as e:
                        logger.error(
                            f"Failed at parsing manifest JSON file {jsonf[0]} with error: {e}"
                        )

            if app_versions:
                apps[appname] = app_versions

        return apps

    def _get_app_files(self, manifest: dict, appdir: Optional[str] = None) -> dict:
        """Get files associated with an app.

        Parameters
        ----------
        manifest : dict
            The app manifest
        appdir : str, optional
            The app directory

        Returns
        -------
        dict
            Dictionary of filename -> _File objects
        """
        source_dir = None
        if appdir is not None:
            source_dir = os.path.join(appdir, "files")

        # Local backend - scan the filesystem
        files: Dict[str, _File] = {}

        if source_dir is None:
            return files

        for ff in glob(os.path.join(source_dir, "**", "*"), recursive=True):
            fname = os.path.relpath(ff, source_dir)
            abspath = os.path.abspath(ff)
            for dir_name in self._get_parent_dirs(fname):
                if dir_name not in files:
                    files[dir_name] = _File(
                        dir_name, os.path.join(source_dir, dir_name)
                    )
            files[fname] = _File(fname, abspath)

        return files

    def _get_parent_dirs(self, name: str):
        """Yield parent directories for a file path."""
        from pathlib import Path

        name_p = Path(name)

        for i in range(1, len(name_p.parts)):
            yield str(Path(*name_p.parts[:i]))

    def _check_folder_validity(self, app_d: str) -> bool:
        """Check if folder path has valid characters."""
        import re

        folder_path = os.path.relpath(app_d, self.root)

        # Check if the folder consists only of lowercase letters, numbers, underscores or path separators
        if not re.match(r"^[a-z0-9_/\\]+$", folder_path):
            logger.warning(
                f'Path {app_d} has invalid characters in the part: "{folder_path}". '
                "Only lowercase letters, numbers and underscores are allowed. "
                "Please fix the path. Skipping..."
            )
            return False
        return True


class _LocalExecutionBackend:
    """Backend for executing apps locally."""

    def __init__(self):
        """Initialize the local execution backend."""

    def prepare_inputs(
        self,
        write_dir: str,
        inputs_dir: str,
        arguments: dict,
        manifest: dict,
        app_files: dict,
        function: str,
        slpm_path: str,
    ) -> None:
        """Prepare inputs for local execution.

        This copies/symlinks input files to the inputs directory and
        writes the inputs.json file.

        Parameters
        ----------
        write_dir : str
            Output directory
        inputs_dir : str
            Directory to write inputs to
        arguments : dict
            Function arguments
        manifest : dict
            Function manifest
        app_files : dict
            App files dictionary
        function : str
            Function name
        slpm_path : str
            The slpm path
        """
        config = _get_config()
        os.makedirs(inputs_dir, exist_ok=True)

        if "outputs" in manifest:
            with open(os.path.join(inputs_dir, "expected_outputs.json"), "w") as f:
                json.dump(manifest["outputs"], f)

        # Track original paths for deduplication
        original_paths = {}

        # Validate and process arguments
        for arg in manifest["params"]:
            name = arg["name"]
            argtype = arg["type"]
            nargs = arg.get("nargs")

            vals = arguments.get(name)
            if vals is None:
                continue

            # Validate type
            vals = _validate_argument_type(name, vals, argtype, nargs)

            # Copy Path-type arguments to folder
            if argtype == "Path" and name not in ("outdir", "scratchdir", "execdir"):
                newvals = self._process_path_arguments(
                    vals,
                    app_files,
                    inputs_dir,
                    write_dir,
                    original_paths,
                    config.symlink,
                )
                arguments[name] = self._normalize_values(newvals)

            # Handle KWARGS type (for HTMD app)
            if argtype == "KWARGS" and vals[0] is not None:
                self._process_kwargs_arguments(
                    vals[0], inputs_dir, write_dir, original_paths, config.symlink
                )
                arguments[name] = vals[0]

        # Write inputs.json
        with open(os.path.join(inputs_dir, "inputs.json"), "w") as f:
            json.dump(
                {"function": function, "slpm_path": slpm_path, "arguments": arguments},
                f,
                indent=4,
            )

        # Write original_paths.json
        with open(os.path.join(inputs_dir, "original_paths.json"), "w") as f:
            json.dump(original_paths, f, indent=4)

    def _process_path_arguments(
        self,
        vals: list,
        app_files: dict,
        inputs_dir: str,
        write_dir: str,
        original_paths: dict,
        use_symlink: bool,
    ) -> list:
        """Process Path-type arguments, copying files to inputs directory."""
        newvals = []

        for val in vals:
            if val is None or (isinstance(val, str) and val == ""):
                continue

            # Handle app://files URIs
            if isinstance(val, str) and val.startswith("app://files"):
                val = app_files[val.replace("app://files/", "")]

            if isinstance(val, _File):
                newvals.append(val.path)
                continue  # Don't copy artifacts

            val = os.path.abspath(val)

            # Deduplicate file names
            outname = os.path.join(inputs_dir, os.path.basename(val))
            if os.path.exists(outname) and val != original_paths.get(outname):
                i = 0
                while os.path.exists(outname):
                    parts = os.path.splitext(os.path.basename(val))
                    outname = os.path.join(inputs_dir, f"{parts[0]}_{i}{parts[1]}")
                    i += 1

            original_paths[outname] = val

            # Copy or symlink
            if use_symlink:
                os.symlink(val, outname)
            else:
                if os.path.isdir(val):
                    shutil.copytree(val, outname)
                else:
                    shutil.copy(val, outname)

            newvals.append(os.path.relpath(outname, write_dir))

        return newvals

    def _process_kwargs_arguments(
        self,
        kwargs_dict: dict,
        inputs_dir: str,
        write_dir: str,
        original_paths: dict,
        use_symlink: bool,
    ) -> None:
        """Process KWARGS type arguments (for HTMD app)."""
        for key, item_vals in kwargs_dict.items():
            if not isinstance(item_vals, (list, tuple)):
                item_vals = [item_vals]

            newvals = []
            for val in item_vals:
                if val is None:
                    continue

                if isinstance(val, _File):
                    newvals.append(val.path)
                    continue

                if not isinstance(val, str) or not os.path.exists(val):
                    newvals.append(val)
                    continue

                val = os.path.abspath(val)

                outname = os.path.join(inputs_dir, os.path.basename(val))
                if os.path.exists(outname) and val != original_paths.get(outname):
                    i = 0
                    while os.path.exists(outname):
                        parts = os.path.splitext(os.path.basename(val))
                        outname = os.path.join(inputs_dir, f"{parts[0]}_{i}{parts[1]}")
                        i += 1

                original_paths[outname] = val

                if use_symlink:
                    os.symlink(val, outname)
                else:
                    if os.path.isdir(val):
                        shutil.copytree(val, outname)
                    else:
                        shutil.copy(val, outname)
                newvals.append(os.path.relpath(outname, write_dir))

            kwargs_dict[key] = self._normalize_values(newvals)

    def _normalize_values(self, vals: list) -> Any:
        """Normalize a list of values to single value or list."""
        if len(vals) == 0:
            return None
        elif len(vals) == 1:
            return vals[0]
        return vals

    def run(
        self,
        dirname: str,
        runsh: str,
        verbose: bool = True,
        **kwargs,
    ) -> None:
        """Run the app locally.

        Parameters
        ----------
        dirname : str
            Directory containing the run script
        runsh : str
            Run script filename
        verbose : bool
            Whether to print output
        **kwargs
            Ignored for local execution
        """
        import subprocess

        logfile_path = os.path.join(dirname, runsh.replace(".sh", ".log"))

        with open(logfile_path, "w") as logfile:
            process = subprocess.Popen(
                ["bash", runsh],
                cwd=dirname,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )

            try:
                for line in process.stdout:
                    if verbose:
                        print(line, end="")
                    logfile.write(line)
                process.wait()
                if process.returncode != 0:
                    raise RuntimeError(
                        f"Job failed with exit code {process.returncode}"
                    )
            except KeyboardInterrupt:
                print("\nKeyboard interrupt detected. Stopping the process...")

                # Try to terminate the process gracefully first
                # This will trigger the trap in the bash script which will clean up the Docker container
                process.terminate()

                # Give it 5 seconds to terminate
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    # If it doesn't terminate, kill it
                    print("Process did not terminate gracefully. Killing it...")
                    process.kill()
                    process.wait()

                raise KeyboardInterrupt("Process interrupted by user")

    def get_status(
        self,
        dirname: str,
        inputs_dir: str,
        slurmq: Any = None,
    ) -> JobStatus:
        """Get job status for local execution.

        Parameters
        ----------
        dirname : str
            Output directory
        inputs_dir : str
            Directory containing inputs
        slurmq : SlurmQueue, optional
            SLURM queue object if submitted to SLURM

        Returns
        -------
        JobStatus
            Current job status
        """
        # Check for completion/error sentinel files
        if os.path.exists(os.path.join(inputs_dir, ".pm.done")) or os.path.exists(
            os.path.join(dirname, ".pm.done")
        ):
            return JobStatus.COMPLETED
        elif os.path.exists(os.path.join(inputs_dir, ".pm.err")) or os.path.exists(
            os.path.join(dirname, ".pm.err")
        ):
            return JobStatus.ERROR

        # Check heartbeat
        heartbeat = os.path.join(inputs_dir, ".pm.alive")
        if not os.path.exists(heartbeat):
            heartbeat = os.path.join(dirname, ".pm.alive")

        if os.path.exists(heartbeat):
            with open(heartbeat, "r") as f:
                timestamp_str = f.read().strip()
                timestamp = None

                if len(timestamp_str):
                    try:
                        timestamp = datetime.datetime.fromisoformat(timestamp_str)
                    except Exception:
                        logger.error(f"Malformed timestamp in {heartbeat}")

                if timestamp is not None:
                    diff = datetime.datetime.now() - timestamp
                    if diff > datetime.timedelta(seconds=JOB_TIMEOUT):
                        return JobStatus.ERROR
                    else:
                        return JobStatus.RUNNING

        # Check SLURM status
        if slurmq is not None:
            from playmolecule._backends._slurm import _get_slurm_status

            return _get_slurm_status(slurmq)

        # Check expected outputs
        outputs_file = os.path.join(inputs_dir, "expected_outputs.json")
        if os.path.exists(outputs_file):
            with open(outputs_file, "r") as f:
                outputs = json.load(f)

            for outf in outputs:
                if len(glob(os.path.join(dirname, outf))) == 0:
                    return JobStatus.RUNNING
            else:
                return JobStatus.COMPLETED

        logger.warning(
            f"Could not yet determine job status for directory {dirname}. "
            "The job might have not started running yet."
        )
        return JobStatus.WAITING_INFO
