#!/usr/bin/env python
# -*- coding: utf-8; -*-

# Copyright (c) 2021, 2022 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
"""This is a driver script auto-generated by Oracle ADS to run Python script in OCI Data Science Jobs.

The following environment variables are used:
CODE_ENTRYPOINT:
    Relative path to the entry script/module file in the user code directory.
ENTRY_FUNCTION:
    Optional, function name in the entry script/module to be invoked.
    If this is not specified, the entry script will be run as Python script.
PYTHON_PATH:
    Optional, additional paths to be added to sys.path for looking up modules and packages.
    The root of the user code will be added by default.
    Multiple paths can be separated by os.pathsep, which is colon(:) for Linux and Mac, semicolon(;) for Windows.
OUTPUT_DIR:
    Optional, output directory to be copied to object storage.
OUTPUT_URI:
    Optional, object storage URI for saving files from the output directory.
WORKING_DIR:
    Optional, the working directory of the user code.
    This can be specified as a relative path from /home/datascience/decompressed_artifact/code
    If this is not specified, the working directory will be /home/datascience/decompressed_artifact/code

This module requires the following package:
oci


"""
import importlib
import json
import logging
import os
import runpy
import subprocess
import sys
from typing import Optional
from urllib.parse import urlparse

import oci


logger = logging.getLogger(__name__)


# Constants
CONST_METADATA = "METADATA"
USER_CODE_DIR = "code"


class OCIHelper:
    @staticmethod
    def init_oci_client(client_class):
        """Initializes OCI client with API key or Resource Principal.

        Parameters
        ----------
        client_class :
            The class of OCI client to be initialized.
        """
        if os.environ.get("OCI_RESOURCE_PRINCIPAL_VERSION"):
            logger.info(
                "Initializing %s with Resource Principal...", client_class.__name__
            )
            client = client_class(
                {}, signer=oci.auth.signers.get_resource_principals_signer()
            )
        else:
            logger.info("Initializing %s with API Key...", {client_class.__name__})
            client = client_class(oci.config.from_file())
        return client

    @staticmethod
    def copy_outputs(output_dir: str, output_uri: dict) -> None:
        """Copies the output files to object storage bucket.

        Parameters
        ----------
        output_dir : str
            Path of the output directory containing files to be copied.
        output_uri : str
            URI of the object storage path to store the output files.
        """
        output_dir = os.path.abspath(os.path.expanduser(output_dir))
        if not os.path.exists(output_dir):
            logger.error("Output directory %s not found.", output_dir)
            return
        if not output_uri:
            logger.error(
                "OUTPUT_URI is not defined in environment variable. No file is copied."
            )
            return
        logger.debug("Copying files in %s to %s...", output_dir, output_uri)
        parsed = urlparse(output_uri)
        bucket_name = parsed.username
        namespace = parsed.hostname
        prefix = parsed.path
        oci_os_client = OCIHelper.init_oci_client(
            oci.object_storage.ObjectStorageClient
        )

        if not prefix:
            prefix = ""
        prefix = prefix.strip("/")

        for path, _, files in os.walk(output_dir):
            for name in files:
                file_path = os.path.join(path, name)

                with open(file_path, "rb") as pkf:
                    # Get the relative path of the file to keep the directory structure
                    relative_path = os.path.relpath(file_path, output_dir)
                    if prefix:
                        file_prefix = os.path.join(prefix, relative_path)
                    else:
                        # Save file to bucket root if prefix is empty.
                        file_prefix = relative_path

                    logger.debug(
                        f"Saving {relative_path} to {bucket_name}@{namespace}/{file_prefix}"
                    )

                    oci_os_client.put_object(
                        namespace,
                        bucket_name,
                        file_prefix,
                        pkf,
                    )


class JobRunner:
    """Contains methods for running the job."""

    def __init__(self, code_dir: str) -> None:
        """Initialize the job runner

        Parameters
        ----------
        code_dir : str
            the path to the directory containing the user code.
        """
        self.code_dir = code_dir

    def setup_python_path(self, python_paths: str):
        """Adds additional python paths.

        Parameters
        ----------
        python_paths : str
            Additional python paths to be added to sys.path.
            Multiple paths can be separated by os.pathsep, which is colon(:) for Linux and Mac.

        """
        python_paths = python_paths.split(os.pathsep)
        python_paths.append(self.code_dir)
        for path in python_paths:
            python_path = os.path.abspath(os.path.expanduser(path))
            if python_path not in sys.path:
                sys.path.append(python_path)
        logger.debug("Python Path: %s", sys.path)
        return self

    def _run_function(self, module_path: str, entry_function: str, argv: list):
        """Runs the entry function in module specified by module path.

        Parameters
        ----------
        module_path : str
            The path to the module containing the entry function.
        entry_function : str
            The name of the entry function.
        argv : list
            Argument list from command line.
        This list will be parsed into positional arguments and keyword arguments.
        """
        spec = importlib.util.spec_from_file_location("module", module_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        method = getattr(module, entry_function)

        args, kwargs = ArgumentParser(argv).parse()

        logger.debug(
            "Invoking method: %s with args: %s, kwargs: %s",
            entry_function,
            args,
            kwargs,
        )
        method(*args, **kwargs)

    def run(
        self,
        entrypoint: str,
        entry_function: Optional[str] = None,
        argv: Optional[list] = None,
    ):
        """Runs the user code

        Parameters
        ----------
        entrypoint : str
            Path to the file serve as the entrypoint.
        entry_function : str, optional
            Name of the function in the entrypoint, by default None.
            If this is not specified, the entrypoint will be run as a python script.
        argv : list, optional
            Arguments to be parsed and passed into the entry_function, by default None.
            The arguments here are used only when running the entry_function.
            When entry_function is None, sys.argv[1:] will be used as command line argument.

        """
        if not argv:
            argv = []

        module_path = entrypoint

        if entry_function:
            logger.info(f"Running function: {entry_function} in {module_path}...")
            self._run_function(module_path, entry_function, argv)
        else:
            logger.info(f"Running script: {module_path}")
            runpy.run_path(module_path, run_name="__main__")
        return self


class ArgumentParser:
    """Contains methods for parsing arguments for entry function."""

    def __init__(self, argument_list: list) -> None:
        """Initialize the parser with a list of arguments

        Parameters
        ----------
        argument_list : list
            A list of arguments.
        """
        self.argument_list = argument_list

    @staticmethod
    def decode_arg(val):
        """Decodes the value of the argument if it is a JSON payload.

        Parameters
        ----------
        val : str
            The argument value in a string.

        Returns
        -------
        Any
            None, if the val is None.
            String value, if the val is a string but not a JSON payload.
            Otherwise, the object after JSON decoded.
        """
        if val is None:
            return None
        try:
            return json.loads(val)
        except json.decoder.JSONDecodeError:
            return val

    @staticmethod
    def join_values(value_list):
        """Joins the values of a keyword argument.

        Parameters
        ----------
        value_list : list
            Values in a list of strings.

        Returns
        -------
        str or None
            The value of the argument as a string.
        """
        if value_list:
            return " ".join(value_list)
        return None

    def parse(self):
        """Parses the arguments

        Returns
        -------
        (list, dict)
            A tuple of positional arguments (list) and keyword arguments (dict).
        """
        args = []
        kwargs = {}
        parsing_kwargs = False
        key = None
        val = []
        for arg in self.argument_list:
            arg = str(arg)
            if len(arg) > 2 and arg.startswith("--"):
                if key:
                    # Save previous key and val
                    kwargs[key] = self.join_values(val)
                parsing_kwargs = True
                key = arg[2:]
                # Reset val
                val = []
            elif parsing_kwargs:
                val.append(arg)
            else:
                args.append(arg)
        # Save the last key and val
        if key:
            kwargs[key] = self.join_values(val)

        args = [self.decode_arg(arg) for arg in args]
        kwargs = {k: self.decode_arg(v) for k, v in kwargs.items()}
        return args, kwargs


def main():
    """The main function for running the job."""
    job_run_ocid = os.environ.get("JOB_RUN_OCID")
    logger.info("Job Run ID is: %s", job_run_ocid)

    # Conda unpack
    try:
        # Ignore the error as conda-unpack may not exists
        subprocess.check_output("conda-unpack", stderr=subprocess.STDOUT, shell=True)
    except subprocess.CalledProcessError as ex:
        logger.debug("conda-unpack exits with non-zero return code %s", ex.returncode)
        logger.debug(ex.output)

    code_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), USER_CODE_DIR)
    # Change the working diretory to the root of user code
    if os.environ.get("WORKING_DIR"):
        working_dir = os.path.join(code_dir, os.environ.get("WORKING_DIR"))
    else:
        working_dir = code_dir
    os.chdir(working_dir)
    # Add working dir to sys.path
    if working_dir not in sys.path:
        sys.path.append(working_dir)

    # Run user code
    runner = JobRunner(code_dir)
    runner.setup_python_path(os.environ.get("PYTHON_PATH", "")).run(
        entrypoint=os.environ.get("CODE_ENTRYPOINT"),
        entry_function=os.environ.get("ENTRY_FUNCTION"),
        argv=sys.argv[1:],
    )

    logger.info("Job completed.")

    # Copy outputs
    if "OUTPUT_DIR" in os.environ:
        logger.debug("OUTPUT_DIR is: %s", os.environ["OUTPUT_DIR"])
        OCIHelper.copy_outputs(
            output_dir=os.environ["OUTPUT_DIR"],
            output_uri=os.environ.get("OUTPUT_URI"),
        )
    else:
        logger.debug("OUTPUT_DIR is not configured. Skipping copy artifacts")


if __name__ == "__main__":
    main()
