import os
import sys
import platform

from metaflow import R, current
from metaflow.metadata import MetaDatum
from metaflow.metadata.util import sync_local_metadata_to_datastore
from metaflow.sidecar import Sidecar
from metaflow.decorators import StepDecorator
from metaflow.exception import MetaflowException
from metaflow.metaflow_config import (
    DEFAULT_CONTAINER_IMAGE,
    DEFAULT_CONTAINER_REGISTRY,
    SNOWPARK_ACCOUNT,
    SNOWPARK_USER,
    SNOWPARK_PASSWORD,
    SNOWPARK_ROLE,
    SNOWPARK_DATABASE,
    SNOWPARK_WAREHOUSE,
    SNOWPARK_SCHEMA,
)

from metaflow.metaflow_config import (
    DATASTORE_LOCAL_DIR,
)

from .snowpark_exceptions import SnowflakeException
from metaflow.plugins.aws.aws_utils import get_docker_registry


class Snowflake(object):
    def __init__(self, connection_params):
        self.connection_params = connection_params

    def session(self):
        # if using the pypi/conda decorator with @snowpark of any step,
        # make sure to pass {'snowflake': '0.11.0'} as well to that step
        try:
            from snowflake.snowpark import Session

            session = Session.builder.configs(self.connection_params).create()
            return session
        except (NameError, ImportError, ModuleNotFoundError):
            raise SnowflakeException(
                "Could not import module 'snowflake'.\n\nInstall Snowflake "
                "Python package (https://pypi.org/project/snowflake/) first.\n"
                "You can install the module by using the @pypi decorator - "
                "Eg: @pypi(packages={'snowflake': '0.11.0'})\n"
            )


class SnowparkDecorator(StepDecorator):
    name = "snowpark"

    defaults = {
        "account": None,
        "user": None,
        "password": None,
        "role": None,
        "database": None,
        "warehouse": None,
        "schema": None,
        "image": None,
        "stage": None,
        "compute_pool": None,
        "volume_mounts": None,
        "external_integration": None,
        "cpu": None,
        "gpu": None,
        "memory": None,
    }

    package_url = None
    package_sha = None
    run_time_limit = None

    def __init__(self, attributes=None, statically_defined=False):
        super(SnowparkDecorator, self).__init__(attributes, statically_defined)

        if not self.attributes["account"]:
            self.attributes["account"] = SNOWPARK_ACCOUNT
        if not self.attributes["user"]:
            self.attributes["user"] = SNOWPARK_USER
        if not self.attributes["password"]:
            self.attributes["password"] = SNOWPARK_PASSWORD
        if not self.attributes["role"]:
            self.attributes["role"] = SNOWPARK_ROLE
        if not self.attributes["database"]:
            self.attributes["database"] = SNOWPARK_DATABASE
        if not self.attributes["warehouse"]:
            self.attributes["warehouse"] = SNOWPARK_WAREHOUSE
        if not self.attributes["schema"]:
            self.attributes["schema"] = SNOWPARK_SCHEMA

        # If no docker image is explicitly specified, impute a default image.
        if not self.attributes["image"]:
            # If metaflow-config specifies a docker image, just use that.
            if DEFAULT_CONTAINER_IMAGE:
                self.attributes["image"] = DEFAULT_CONTAINER_IMAGE
            # If metaflow-config doesn't specify a docker image, assign a
            # default docker image.
            else:
                # Metaflow-R has its own default docker image (rocker family)
                if R.use_r():
                    self.attributes["image"] = R.container_image()
                # Default to vanilla Python image corresponding to major.minor
                # version of the Python interpreter launching the flow.
                self.attributes["image"] = "python:%s.%s" % (
                    platform.python_version_tuple()[0],
                    platform.python_version_tuple()[1],
                )

        # Assign docker registry URL for the image.
        if not get_docker_registry(self.attributes["image"]):
            if DEFAULT_CONTAINER_REGISTRY:
                self.attributes["image"] = "%s/%s" % (
                    DEFAULT_CONTAINER_REGISTRY.rstrip("/"),
                    self.attributes["image"],
                )

    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
    # to understand where these functions are invoked in the lifecycle of a
    # Metaflow flow.
    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
        # Set internal state.
        self.logger = logger
        self.environment = environment
        self.step = step
        self.flow_datastore = flow_datastore

        if any([deco.name == "parallel" for deco in decos]):
            raise MetaflowException(
                "Step *{step}* contains a @parallel decorator "
                "with the @snowpark decorator. @parallel is not supported with @snowpark.".format(
                    step=step
                )
            )

    def package_init(self, flow, step_name, environment):
        try:
            # Snowflake is a soft dependency.
            from snowflake.snowpark import Session
        except (NameError, ImportError, ModuleNotFoundError):
            raise SnowflakeException(
                "Could not import module 'snowflake'.\n\nInstall Snowflake "
                "Python package (https://pypi.org/project/snowflake/) first.\n"
                "You can install the module by executing - "
                "%s -m pip install snowflake\n"
                "or equivalent through your favorite Python package manager."
                % sys.executable
            )

    def runtime_init(self, flow, graph, package, run_id):
        # Set some more internal state.
        self.flow = flow
        self.graph = graph
        self.package = package
        self.run_id = run_id

    def runtime_task_created(
        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
    ):
        if not is_cloned:
            self._save_package_once(self.flow_datastore, self.package)

    def runtime_step_cli(
        self, cli_args, retry_count, max_user_code_retries, ubf_context
    ):
        if retry_count <= max_user_code_retries:
            cli_args.commands = ["snowpark", "step"]
            cli_args.command_args.append(self.package_sha)
            cli_args.command_args.append(self.package_url)
            cli_args.command_options.update(self.attributes)
            cli_args.command_options["run-time-limit"] = self.run_time_limit
            if not R.use_r():
                cli_args.entrypoint[0] = sys.executable

    def task_pre_step(
        self,
        step_name,
        task_datastore,
        metadata,
        run_id,
        task_id,
        flow,
        graph,
        retry_count,
        max_retries,
        ubf_context,
        inputs,
    ):
        self.metadata = metadata
        self.task_datastore = task_datastore

        # this path will exist within snowpark container services
        login_token = open("/snowflake/session/token", "r").read()
        connection_params = {
            "account": os.environ.get("SNOWFLAKE_ACCOUNT"),
            "host": os.environ.get("SNOWFLAKE_HOST"),
            "authenticator": "oauth",
            "token": login_token,
            "database": os.environ.get("SNOWFLAKE_DATABASE"),
            "schema": os.environ.get("SNOWFLAKE_SCHEMA"),
            "autocommit": True,
            "client_session_keep_alive": True,
        }

        # SNOWFLAKE_WAREHOUSE is injected explicitly by us
        # but is not really required. So if it exists, we use it in
        # connection parameters
        if os.environ.get("SNOWFLAKE_WAREHOUSE"):
            connection_params["warehouse"] = os.environ.get("SNOWFLAKE_WAREHOUSE")

        snowflake_obj = Snowflake(connection_params)
        current._update_env({"snowflake": snowflake_obj})

        meta = {}
        if "METAFLOW_SNOWPARK_WORKLOAD" in os.environ:
            meta["snowflake-account"] = os.environ.get("SNOWFLAKE_ACCOUNT")
            meta["snowflake-database"] = os.environ.get("SNOWFLAKE_DATABASE")
            meta["snowflake-schema"] = os.environ.get("SNOWFLAKE_SCHEMA")
            meta["snowflake-host"] = os.environ.get("SNOWFLAKE_HOST")
            meta["snowflake-service-name"] = os.environ.get("SNOWFLAKE_SERVICE_NAME")

            self._save_logs_sidecar = Sidecar("save_logs_periodically")
            self._save_logs_sidecar.start()

        if len(meta) > 0:
            entries = [
                MetaDatum(
                    field=k,
                    value=v,
                    type=k,
                    tags=["attempt_id:{0}".format(retry_count)],
                )
                for k, v in meta.items()
                if v is not None
            ]
            # Register book-keeping metadata for debugging.
            metadata.register_metadata(run_id, step_name, task_id, entries)

    def task_finished(
        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
    ):
        if "METAFLOW_SNOWPARK_WORKLOAD" in os.environ:
            if hasattr(self, "metadata") and self.metadata.TYPE == "local":
                # Note that the datastore is *always* Amazon S3 (see
                # runtime_task_created function).
                sync_local_metadata_to_datastore(
                    DATASTORE_LOCAL_DIR, self.task_datastore
                )

        try:
            self._save_logs_sidecar.terminate()
        except:
            # Best effort kill
            pass

    @classmethod
    def _save_package_once(cls, flow_datastore, package):
        if cls.package_url is None:
            cls.package_url, cls.package_sha = flow_datastore.save_data(
                [package.blob], len_hint=1
            )[0]
