# (c) 2015-2023 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
import os
from playmolecule._appfiles import _File, _Artifacts
from playmolecule._dynamicfunction import _create_dynamic_function
import stat
import logging

logger = logging.getLogger(__name__)


_function_dict = {}


_JOB_TIMEOUT = 60  # Timeout after which if there is no newer date in .pm.alive we consider the job dead


class KWARGS(dict):
    pass


def _inner_function(
    args: dict,
    function_name: str,
    function_resources: dict,
    func_manifest: dict,
    slpm_path: str,
    files: dict,
    run_sh: str,
):
    from playmolecule._public_api import ExecutableDirectory, JobStatus
    from playmolecule._config import _get_config
    from playmolecule._backends import _get_execution_backend
    from playmolecule._backends._http import _HTTPExecutionBackend
    from datetime import datetime
    import uuid
    import os

    config = _get_config()

    assert (
        "outdir" in args or "execdir" in args
    ), "Functions must accept either 'outdir' or 'execdir'"

    if "outdir" in args:
        outdir = args["outdir"]
        args["outdir"] = "."
    elif "execdir" in args:
        if isinstance(args["execdir"], _File):
            args["execdir"] = args["execdir"].path
        outdir = args["execdir"]
        args["execdir"] = "."

    # Get execution backend
    backend = _get_execution_backend()

    # HTTP backend path
    if isinstance(backend, _HTTPExecutionBackend):
        input_dict = backend.prepare_inputs(args, func_manifest, files, slpm_path)
        ed = ExecutableDirectory(
            dirname=outdir,
            inputs_dir=outdir,
            _execution_resources=function_resources,
            _input_json=input_dict,
            _execution_backend=backend,
        )
        if not config.blocking:
            return ed
        else:
            import time

            ed.run()
            time.sleep(2)
            while ed.status not in (JobStatus.COMPLETED, JobStatus.ERROR):
                time.sleep(3)
            return None

    # Local execution path
    if "scratchdir" in args:
        if args["scratchdir"] is None:
            args["scratchdir"] = os.path.join(outdir, "scratch")
        os.makedirs(args["scratchdir"], exist_ok=True)

    now = datetime.now()
    identifier = f"_{now.strftime(r'%d_%m_%Y_%H_%M')}_{uuid.uuid4().hex[:8]}"
    inputs_dir = os.path.join(outdir, f"run{identifier}")

    # Use backend's prepare_inputs
    backend.prepare_inputs(
        outdir, inputs_dir, args, func_manifest, files, function_name, slpm_path
    )

    target_run_sh = None
    if run_sh is not None and len(run_sh):
        target_run_sh = os.path.join(outdir, f"run{identifier}.sh")
        with open(target_run_sh, "w") as f:
            f.write(run_sh)
        st = os.stat(target_run_sh)
        os.chmod(target_run_sh, st.st_mode | stat.S_IEXEC)

    return ExecutableDirectory(
        dirname=outdir,
        inputs_dir=inputs_dir,
        _execution_resources=function_resources,
        _execution_backend=backend,
    )


def __ensure_module_path(module_path):
    """
    Ensures that all parts of the dotted module path exist in sys.modules
    and returns the final module object.
    """
    import sys
    import types

    parts = module_path.split(".")
    full_path = ""
    parent = None

    for part in parts:
        full_path = f"{full_path}.{part}" if full_path else part
        if full_path not in sys.modules:
            mod = types.ModuleType(full_path)
            sys.modules[full_path] = mod
            if parent:
                setattr(parent, part, mod)
        parent = sys.modules[full_path]

    return sys.modules[full_path]


def _manifest_to_func(appname, app_versions):
    from copy import deepcopy

    for version in app_versions:
        manifest = app_versions[version]["manifest"]
        new_mode = True
        if (
            "functions" not in manifest
        ):  # TODO: Deprecate eventually. Just for backwards compatibility
            new_mode = False
            manifest = deepcopy(manifest)
            try:
                manifest["functions"] = [
                    {
                        "function": (
                            manifest["container_config"]["appfunction"]
                            if "container_config" in manifest
                            else "main"
                        ),
                        "env": "base",
                        "resources": manifest.get("resources", None),
                        "examples": manifest.get("examples", []),
                        "params": manifest["params"],
                        "tests": manifest.get("tests", {}),
                        "outputs": manifest.get("outputs", {}),
                        "description": manifest["description"],
                    }
                ]
            except Exception:
                import traceback

                logger.error(
                    f"Failed to parse manifest for app {appname} version {version} due to error:\n{traceback.format_exc()}"
                )
                continue

        module_path = f"playmolecule.apps.{appname}.{version}"
        module = __ensure_module_path(module_path)

        setattr(module, "__manifest__", deepcopy(manifest))

        run_sh = app_versions[version].get("run.sh")

        app_artifacts = None
        source_dir = app_versions[version].get("appdir")
        if source_dir is not None:
            source_dir = os.path.join(source_dir, "files")
        app_files = app_versions[version].get("files")
        setattr(module, "files", app_files)

        artifact_dict = manifest.get("artifacts", manifest.get("datasets", {}))
        if len(artifact_dict):
            app_artifacts = _Artifacts(artifact_dict, app_files)
            setattr(module, "artifacts", app_artifacts)
            setattr(module, "datasets", app_artifacts)

        func_names = []
        for idx, func_mani in enumerate(manifest["functions"]):
            try:
                func_name = func_mani["function"].split(".")[-1]
                if func_name == "main":
                    func_name = appname

                # Create dynamic function from manifest
                exposes_ports = ""
                _run_sh = run_sh
                if func_mani.get("exposes_ports", None) is not None:
                    exposes_ports = ",".join(func_mani["exposes_ports"])
                if run_sh is not None:
                    _run_sh = run_sh.replace("{exposes_ports}", exposes_ports)

                _func = _create_dynamic_function(
                    func_manifest=func_mani,
                    appname=appname,
                    version=version,
                    run_sh=_run_sh,
                    app_files=app_files,
                    app_artifacts=app_artifacts,
                )
                func_names.append(func_name)
                setattr(module, func_name, _func)
                _function_dict[f"{module_path}.{func_name}"] = _func
            except Exception:
                import traceback

                logger.error(
                    f"Failed to parse manifest for app {appname} version {version} with error: {traceback.format_exc()}"
                )
    return func_names


def _link_latest_version(appname, latest, func_names):
    import sys

    # Link the latest version of the app to the top level module
    parent_module = sys.modules[f"playmolecule.apps.{appname}"]
    latest_module = sys.modules[f"playmolecule.apps.{appname}.{latest}"]
    for symbol in func_names + [
        "artifacts",
        "datasets",
        "files",
        "tests",
        "__manifest__",
    ]:
        if symbol not in parent_module.__dict__ and symbol in latest_module.__dict__:
            setattr(
                sys.modules[f"playmolecule.apps.{appname}"],
                symbol,
                getattr(
                    sys.modules.get(f"playmolecule.apps.{appname}.{latest}"), symbol
                ),
            )


def _find_apps(registries):
    """Set the root directory for apps using the backend registry."""
    from natsort import natsorted
    from playmolecule._backends import _get_manifest_backend

    logger.info(f"Setting PlayMolecule source(s) to {registries}")

    # Get the appropriate manifest backend
    backends = []
    for registry in registries:
        backend = _get_manifest_backend(registry)
        if backend is None:
            raise RuntimeError(f"No backend found for root: {registry}")
        backends.append(backend)

    app_manifests = {}
    for backend in backends:
        new_manifests = backend.get_apps()
        for key in new_manifests:
            if key not in app_manifests:
                app_manifests[key] = new_manifests[key]

    for appname in app_manifests:
        func_names = _manifest_to_func(appname, app_manifests[appname])
        versions = list(app_manifests[appname].keys())
        versions.sort(key=lambda s: natsorted(s))
        _link_latest_version(appname, versions[-1], func_names)
