import copy
import logging
import os
import posixpath
import string
import tempfile
import time
from abc import ABC
from contextlib import contextmanager
from typing import Any, Dict, Generator, Optional, Tuple, cast

import yaml
from typing_extensions import Unpack

from snowflake.ml._internal import file_utils
from snowflake.ml._internal.exceptions import (
    error_codes,
    exceptions as snowml_exceptions,
)
from snowflake.ml._internal.utils import identifier, query_result_checker
from snowflake.ml.model import type_hints
from snowflake.ml.model._deploy_client.image_builds import (
    base_image_builder,
    client_image_builder,
    docker_context,
    server_image_builder,
)
from snowflake.ml.model._deploy_client.snowservice import deploy_options, instance_types
from snowflake.ml.model._deploy_client.utils import (
    constants,
    image_registry_client,
    snowservice_client,
)
from snowflake.ml.model._packager.model_meta import model_meta, model_meta_schema
from snowflake.snowpark import Session

logger = logging.getLogger(__name__)


@contextmanager
def _debug_aware_tmp_directory(debug_dir: Optional[str] = None) -> Generator[str, None, None]:
    """Debug-aware directory provider.

    Args:
        debug_dir: A folder for deploymement context.

    Yields:
        A directory path to write deployment artifacts
    """
    create_temp = False
    if debug_dir:
        directory_path = debug_dir
    else:
        temp_dir_context = tempfile.TemporaryDirectory()
        directory_path = temp_dir_context.name
        create_temp = True
    try:
        yield directory_path
    finally:
        if create_temp:
            temp_dir_context.cleanup()


def _deploy(
    session: Session,
    *,
    model_id: str,
    model_meta: model_meta.ModelMetadata,
    service_func_name: str,
    model_zip_stage_path: str,
    deployment_stage_path: str,
    target_method: str,
    **kwargs: Unpack[type_hints.SnowparkContainerServiceDeployOptions],
) -> type_hints.SnowparkContainerServiceDeployDetails:
    """Entrypoint for model deployment to SnowService. This function will trigger a docker image build followed by
    workflow deployment to SnowService.

    Args:
        session: Snowpark session
        model_id: Unique hex string of length 32, provided by model registry.
        model_meta: Model Metadata.
        service_func_name: The service function name in SnowService associated with the created service.
        model_zip_stage_path: Path to model zip file in stage. Note that this path has a "@" prefix.
        deployment_stage_path: Path to stage containing deployment artifacts.
        target_method: The name of the target method to be deployed.
        **kwargs: various SnowService deployment options.

    Returns:
        Deployment details for SPCS.

    Raises:
        SnowflakeMLException: Raised when model_id is empty.
        SnowflakeMLException: Raised when service_func_name is empty.
        SnowflakeMLException: Raised when model_stage_file_path is empty.
    """
    snowpark_logger = logging.getLogger("snowflake.snowpark")
    snowflake_connector_logger = logging.getLogger("snowflake.connector")
    snowpark_log_level = snowpark_logger.level
    snowflake_connector_log_level = snowflake_connector_logger.level
    try:
        # Setting appropriate log level to prevent console from being polluted by vast amount of snowpark and snowflake
        # connector logging.
        snowpark_logger.setLevel(logging.WARNING)
        snowflake_connector_logger.setLevel(logging.WARNING)
        if not model_id:
            raise snowml_exceptions.SnowflakeMLException(
                error_code=error_codes.INVALID_ARGUMENT,
                original_exception=ValueError(
                    'Must provide a non-empty string for "model_id" when deploying to Snowpark Container Services'
                ),
            )
        if not service_func_name:
            raise snowml_exceptions.SnowflakeMLException(
                error_code=error_codes.INVALID_ARGUMENT,
                original_exception=ValueError(
                    'Must provide a non-empty string for "service_func_name"'
                    " when deploying to Snowpark Container Services"
                ),
            )
        if not model_zip_stage_path:
            raise snowml_exceptions.SnowflakeMLException(
                error_code=error_codes.INVALID_ARGUMENT,
                original_exception=ValueError(
                    'Must provide a non-empty string for "model_stage_file_path"'
                    " when deploying to Snowpark Container Services"
                ),
            )
        if not deployment_stage_path:
            raise snowml_exceptions.SnowflakeMLException(
                error_code=error_codes.INVALID_ARGUMENT,
                original_exception=ValueError(
                    'Must provide a non-empty string for "deployment_stage_path"'
                    " when deploying to Snowpark Container Services"
                ),
            )

        # Remove full qualified name to avoid double quotes corrupting the service spec
        model_zip_stage_path = model_zip_stage_path.replace('"', "")
        deployment_stage_path = deployment_stage_path.replace('"', "")

        assert model_zip_stage_path.startswith("@"), f"stage path should start with @, actual: {model_zip_stage_path}"
        assert deployment_stage_path.startswith("@"), f"stage path should start with @, actual: {deployment_stage_path}"
        options = deploy_options.SnowServiceDeployOptions.from_dict(cast(Dict[str, Any], kwargs))

        model_meta_deploy = copy.deepcopy(model_meta)
        # Set conda-forge as backup channel for SPCS deployment
        if "conda-forge" not in model_meta_deploy.env._conda_dependencies:
            model_meta_deploy.env._conda_dependencies["conda-forge"] = []
        if options.use_gpu:
            # Make mypy happy
            assert options.num_gpus is not None
            if model_meta_deploy.env.cuda_version is None:
                raise snowml_exceptions.SnowflakeMLException(
                    error_code=error_codes.INVALID_ARGUMENT,
                    original_exception=ValueError(
                        "You are requesting GPUs for models that do not use a GPU or does not have CUDA version set."
                    ),
                )
            if model_meta.env.cuda_version:
                model_meta_deploy.env.generate_env_for_cuda()
        else:
            # If user does not need GPU, we set this copies cuda_version to None, thus when Image builder gets a
            # not-None cuda_version, it gets to know that GPU is required.
            model_meta_deploy.env._cuda_version = None

        _validate_compute_pool(session, options=options)

        # TODO[shchen]: SNOW-863701, Explore ways to prevent entire model zip being downloaded during deploy step
        #  (for both warehouse and snowservice deployment)
        # One alternative is for model registry to duplicate the model metadata and env dependency storage from model
        # zip so that we don't have to pull down the entire model zip.
        ss_deployment = SnowServiceDeployment(
            session=session,
            model_id=model_id,
            model_meta=model_meta_deploy,
            service_func_name=service_func_name,
            model_zip_stage_path=model_zip_stage_path,  # Pass down model_zip_stage_path for service spec file
            deployment_stage_path=deployment_stage_path,
            target_method=target_method,
            options=options,
        )
        return ss_deployment.deploy()
    finally:
        # Preserve the original logging level.
        snowpark_logger.setLevel(snowpark_log_level)
        snowflake_connector_logger.setLevel(snowflake_connector_log_level)


def _validate_compute_pool(session: Session, *, options: deploy_options.SnowServiceDeployOptions) -> None:
    # Remove full qualified name to avoid double quotes, which does not work well in desc compute pool syntax.
    compute_pool = options.compute_pool.replace('"', "")
    sql = f"DESC COMPUTE POOL {compute_pool}"
    result = (
        query_result_checker.SqlResultValidator(
            session=session,
            query=sql,
        )
        .has_column("instance_family")
        .has_column("state")
        .has_column("auto_resume")
        .has_dimensions(expected_rows=1)
        .validate()
    )

    state = result[0]["state"]
    auto_resume = bool(result[0]["auto_resume"])

    if state == "SUSPENDED":
        if not auto_resume:
            raise snowml_exceptions.SnowflakeMLException(
                error_code=error_codes.INVALID_SNOWPARK_COMPUTE_POOL,
                original_exception=RuntimeError(
                    "The compute pool you are requesting to use is suspended without auto-resume enabled"
                ),
            )

    elif state not in ["ACTIVE", "IDLE"]:
        raise snowml_exceptions.SnowflakeMLException(
            error_code=error_codes.INVALID_SNOWPARK_COMPUTE_POOL,
            original_exception=RuntimeError(
                "The compute pool you are requesting to use is not in the ACTIVE/IDLE status."
            ),
        )

    if options.use_gpu:
        assert options.num_gpus is not None
        request_gpus = options.num_gpus
        instance_family = result[0]["instance_family"]
        if instance_family in instance_types.INSTANCE_TYPE_TO_GPU_COUNT:
            gpu_capacity = instance_types.INSTANCE_TYPE_TO_GPU_COUNT[instance_family]
            if request_gpus > gpu_capacity:
                raise snowml_exceptions.SnowflakeMLException(
                    error_code=error_codes.INVALID_SNOWPARK_COMPUTE_POOL,
                    original_exception=RuntimeError(
                        f"GPU request exceeds instance capability; {instance_family} instance type has total "
                        f"capacity of {gpu_capacity} GPU, yet a request was made for {request_gpus} GPUs."
                    ),
                )
        else:
            logger.warning(f"Unknown instance type: {instance_family}, skipping GPU validation")


def _get_or_create_image_repo(session: Session, *, service_func_name: str, image_repo: Optional[str] = None) -> str:
    def _sanitize_dns_url(url: str) -> str:
        # Align with existing SnowService image registry url standard.
        return url.lower()

    if image_repo:
        return _sanitize_dns_url(image_repo)

    try:
        conn = session._conn._conn
        # We try to use the same db and schema as the service function locates, as we could retrieve those information
        # if that is a fully qualified one. If not we use the current session one.
        (_db, _schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)
        db = _db if _db is not None else conn._database
        schema = _schema if _schema is not None else conn._schema
        assert isinstance(db, str) and isinstance(schema, str)

        client = snowservice_client.SnowServiceClient(session)
        client.create_image_repo(identifier.get_schema_level_object_identifier(db, schema, constants.SNOWML_IMAGE_REPO))
        sql = f"SHOW IMAGE REPOSITORIES LIKE '{constants.SNOWML_IMAGE_REPO}' IN SCHEMA {'.'.join([db, schema])}"
        result = (
            query_result_checker.SqlResultValidator(
                session=session,
                query=sql,
            )
            .has_column("repository_url")
            .has_dimensions(expected_rows=1)
            .validate()
        )
        repository_url = result[0]["repository_url"]
        return str(repository_url)
    except Exception as e:
        raise snowml_exceptions.SnowflakeMLException(
            error_code=error_codes.INTERNAL_SNOWPARK_CONTAINER_SERVICE_ERROR,
            original_exception=RuntimeError("Failed to retrieve image repo URL"),
        ) from e


class SnowServiceDeployment(ABC):
    """
    Class implementation that encapsulates image build and workflow deployment to SnowService
    """

    def __init__(
        self,
        session: Session,
        model_id: str,
        model_meta: model_meta.ModelMetadata,
        service_func_name: str,
        model_zip_stage_path: str,
        deployment_stage_path: str,
        target_method: str,
        options: deploy_options.SnowServiceDeployOptions,
    ) -> None:
        """Initialization

        Args:
            session: Snowpark session
            model_id: Unique hex string of length 32, provided by model registry; if not provided, auto-generate one for
                        resource naming.The model_id serves as an idempotent key throughout the deployment workflow.
            model_meta: Model Metadata.
            service_func_name: The service function name in SnowService associated with the created service.
            model_zip_stage_path: Path to model zip file in stage.
            deployment_stage_path: Path to stage containing deployment artifacts.
            target_method: The name of the target method to be deployed.
            options: A SnowServiceDeployOptions object containing deployment options.
        """

        self.session = session
        self.id = model_id
        self.model_meta = model_meta
        self.service_func_name = service_func_name
        self.model_zip_stage_path = model_zip_stage_path
        self.options = options
        self.target_method = target_method
        (db, schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)

        self._service_name = identifier.get_schema_level_object_identifier(db, schema, f"service_{model_id}")
        # Spec file and future deployment related artifacts will be stored under {stage}/models/{model_id}
        self._model_artifact_stage_location = posixpath.join(deployment_stage_path, "models", self.id)
        self.debug_dir: Optional[str] = None
        if self.options.debug_mode:
            self.debug_dir = tempfile.mkdtemp()
            logger.warning(f"Debug model is enabled, deployment artifacts will be available in {self.debug_dir}")

    def deploy(self) -> type_hints.SnowparkContainerServiceDeployDetails:
        """
        This function triggers image build followed by workflow deployment to SnowService.

        Returns:
            Deployment details.
        """
        if self.options.prebuilt_snowflake_image:
            logger.warning(f"Skipped image build. Use prebuilt image: {self.options.prebuilt_snowflake_image}")
            full_image_name = self.options.prebuilt_snowflake_image
            (service_spec, service_function_sql) = self._deploy_workflow(self.options.prebuilt_snowflake_image)
        else:
            with _debug_aware_tmp_directory(debug_dir=self.debug_dir) as context_dir:
                extra_kwargs = {}
                if self.options.model_in_image:
                    extra_kwargs = {
                        "session": self.session,
                        "model_zip_stage_path": self.model_zip_stage_path,
                    }
                dc = docker_context.DockerContext(
                    context_dir=context_dir,
                    model_meta=self.model_meta,
                    **extra_kwargs,  # type: ignore[arg-type]
                )
                dc.build()
                image_repo = _get_or_create_image_repo(
                    self.session, service_func_name=self.service_func_name, image_repo=self.options.image_repo
                )
                full_image_name = self._get_full_image_name(image_repo=image_repo, context_dir=context_dir)
                registry_client = image_registry_client.ImageRegistryClient(self.session)

                if not self.options.force_image_build and registry_client.image_exists(full_image_name=full_image_name):
                    logger.warning(
                        f"Similar environment detected. Using existing image {full_image_name} to skip image "
                        f"build. To disable this feature, set 'force_image_build=True' in deployment options"
                    )
                else:
                    logger.warning(
                        "Building the Docker image and deploying to Snowpark Container Service. "
                        "This process may take anywhere from a few minutes to a longer period for GPU-based models."
                    )
                    start = time.time()
                    self._build_and_upload_image(
                        context_dir=context_dir, image_repo=image_repo, full_image_name=full_image_name
                    )
                    end = time.time()
                    logger.info(f"Time taken to build and upload image to registry: {end - start:.2f} seconds")
                    logger.warning(
                        f"Image successfully built! For future model deployments, the image will be reused if "
                        f"possible, saving model deployment time. To enforce using the same image, include "
                        f"'prebuilt_snowflake_image': '{full_image_name}' in the deploy() function's options."
                    )

                # Adding the model name as an additional tag to the existing image, excluding the version to prevent
                # excessive tags and also due to version not available in current model metadata. This will allow
                # users to associate images with specific models and perform relevant image registry actions. In the
                # event that model dependencies change across versions, a new image hash will be computed, resulting in
                # a new image.
                try:
                    registry_client.add_tag_to_remote_image(
                        original_full_image_name=full_image_name, new_tag=self.model_meta.name
                    )
                except Exception as e:
                    # Proceed to the deployment with a warning message.
                    logger.warning(f"Failed to add tag {self.model_meta.name} to image {full_image_name}: {str(e)}")
                (service_spec, service_function_sql) = self._deploy_workflow(full_image_name)
        return type_hints.SnowparkContainerServiceDeployDetails(
            image_name=full_image_name,
            service_spec=service_spec,
            service_function_sql=service_function_sql,
        )

    def _get_full_image_name(self, image_repo: str, context_dir: str) -> str:
        """Return a valid full image name that consists of image name and tag. e.g
        org-account.registry.snowflakecomputing.com/db/schema/repo/image:latest

        Args:
            image_repo: image repo path, e.g. org-account.registry.snowflakecomputing.com/db/schema/repo
            context_dir: the local docker context directory, which consists everything needed to build the docker image.

        Returns:
            Full image name.
        """
        image_repo = _get_or_create_image_repo(
            self.session, service_func_name=self.service_func_name, image_repo=self.options.image_repo
        )

        # We skip "MODEL_METADATA_FILE" as it contains information that will always lead to cache misses.  This isn't an
        # issue because model dependency is also captured in the model env/ folder, which will be hashed. The aim is to
        # reuse the same Docker image even if the user logs a similar model without new dependencies.
        docker_context_dir_hash = file_utils.hash_directory(
            context_dir, ignore_hidden=True, excluded_files=[model_meta.MODEL_METADATA_FILE]
        )
        # By default, we associate a 'latest' tag with each of our created images for easy existence checking.
        # Additional tags are added for readability.
        return f"{image_repo}/{docker_context_dir_hash}:{constants.LATEST_IMAGE_TAG}"

    def _build_and_upload_image(self, context_dir: str, image_repo: str, full_image_name: str) -> None:
        """Handles image build and upload to image registry.

        Args:
            context_dir: the local docker context directory, which consists everything needed to build the docker image.
            image_repo: image repo path, e.g. org-account.registry.snowflakecomputing.com/db/schema/repo
            full_image_name: Full image name consists of image name and image tag.
        """
        image_builder: base_image_builder.ImageBuilder
        if self.options.enable_remote_image_build:
            image_builder = server_image_builder.ServerImageBuilder(
                context_dir=context_dir,
                full_image_name=full_image_name,
                image_repo=image_repo,
                session=self.session,
                artifact_stage_location=self._model_artifact_stage_location,
                compute_pool=self.options.compute_pool,
            )
        else:
            image_builder = client_image_builder.ClientImageBuilder(
                context_dir=context_dir, full_image_name=full_image_name, image_repo=image_repo, session=self.session
            )
        image_builder.build_and_upload_image()

    def _prepare_and_upload_artifacts_to_stage(self, image: str) -> str:
        """Constructs and upload service spec to stage.

        Args:
            image: Name of the image to create SnowService container from.

        Returns:
            Service spec string.
        """
        if self.options.model_in_image:
            spec_template_path = file_utils.resolve_zip_import_path(
                os.path.join(os.path.dirname(__file__), "templates/service_spec_template_with_model")
            )
        else:
            spec_template_path = file_utils.resolve_zip_import_path(
                os.path.join(os.path.dirname(__file__), "templates/service_spec_template")
            )

        with _debug_aware_tmp_directory(self.debug_dir) as dir_path:
            spec_file_path = os.path.join(dir_path, f"{constants.SERVICE_SPEC}.yaml")

            with open(spec_template_path, encoding="utf-8") as template, open(
                spec_file_path, "w+", encoding="utf-8"
            ) as spec_file:
                assert self.model_zip_stage_path.startswith("@")
                norm_stage_path = posixpath.normpath(identifier.remove_prefix(self.model_zip_stage_path, "@"))
                # Ensure model stage path has root prefix as stage mount will it mount it to root.
                absolute_model_stage_path = os.path.join("/", norm_stage_path)
                (db, schema, stage, path) = identifier.parse_schema_level_object_identifier(norm_stage_path)
                substitutes = {
                    "image": image,
                    "predict_endpoint_name": constants.PREDICT,
                    "model_stage": identifier.get_schema_level_object_identifier(db, schema, stage),
                    "model_zip_stage_path": absolute_model_stage_path,
                    "inference_server_container_name": constants.INFERENCE_SERVER_CONTAINER,
                    "target_method": self.target_method,
                    "num_workers": self.options.num_workers,
                    "use_gpu": self.options.use_gpu,
                }
                if self.options.model_in_image:
                    del substitutes["model_stage"]
                    del substitutes["model_zip_stage_path"]
                content = string.Template(template.read()).substitute(substitutes)
                content_dict = yaml.safe_load(content)
                if self.options.use_gpu:
                    container = content_dict["spec"]["container"][0]
                    # TODO[shchen]: SNOW-871538, external dependency that only single GPU is supported on SnowService.
                    # GPU limit has to be specified in order to trigger the workload to be run on GPU in SnowService.
                    container["resources"] = {
                        "limits": {"nvidia.com/gpu": self.options.num_gpus},
                        "requests": {"nvidia.com/gpu": self.options.num_gpus},
                    }

                    # Make LLM use case sequential
                    if any(
                        model_blob_meta.model_type == "huggingface_pipeline" or model_blob_meta.model_type == "llm"
                        for model_blob_meta in self.model_meta.models.values()
                    ):
                        container["env"]["_CONCURRENT_REQUESTS_MAX"] = 1

                yaml.dump(content_dict, spec_file)
                spec_file.seek(0)
                spec_file_yaml_string = spec_file.read()
                logger.debug(f"Create service spec: \n {spec_file_yaml_string}")

            self.session.file.put(
                local_file_name=spec_file_path,
                stage_location=self._model_artifact_stage_location,
                auto_compress=False,
                overwrite=True,
            )
            logger.debug(
                f"Uploaded spec file {os.path.basename(spec_file_path)} " f"to {self._model_artifact_stage_location}"
            )
            return spec_file_yaml_string

    def _get_max_batch_rows(self) -> Optional[int]:
        # To avoid too large batch in HF LLM case
        max_batch_rows = None
        if self.options.use_gpu:
            for model_blob_meta in self.model_meta.models.values():
                batch_size = None
                if model_blob_meta.model_type == "huggingface_pipeline":
                    model_blob_options_hf = cast(
                        model_meta_schema.HuggingFacePipelineModelBlobOptions, model_blob_meta.options
                    )
                    batch_size = model_blob_options_hf["batch_size"]
                if model_blob_meta.model_type == "llm":
                    model_blob_options_llm = cast(model_meta_schema.LLMModelBlobOptions, model_blob_meta.options)
                    batch_size = model_blob_options_llm["batch_size"]
                if batch_size:
                    if max_batch_rows is None:
                        max_batch_rows = batch_size
                    else:
                        max_batch_rows = min(batch_size, max_batch_rows)
        return max_batch_rows

    def _deploy_workflow(self, image: str) -> Tuple[str, str]:
        """This function handles workflow deployment to SnowService with the given image.

        Args:
            image: Name of the image to create SnowService container from.

        Returns:
            Tuple of (service spec, service function sql).
        """

        service_spec_string = self._prepare_and_upload_artifacts_to_stage(image)
        client = snowservice_client.SnowServiceClient(self.session)
        spec_stage_location = posixpath.join(
            self._model_artifact_stage_location.rstrip("/"), f"{constants.SERVICE_SPEC}.yaml"
        )
        client.create_or_replace_service(
            service_name=self._service_name,
            compute_pool=self.options.compute_pool,
            spec_stage_location=spec_stage_location,
            min_instances=self.options.min_instances,
            max_instances=self.options.max_instances,
        )
        logger.info(f"Wait for service {self._service_name} to become ready...")
        client.block_until_resource_is_ready(
            resource_name=self._service_name, resource_type=constants.ResourceType.SERVICE
        )
        logger.info(f"Service {self._service_name} is ready. Creating service function...")

        service_function_sql = client.create_or_replace_service_function(
            service_func_name=self.service_func_name,
            service_name=self._service_name,
            endpoint_name=constants.PREDICT,
            max_batch_rows=self._get_max_batch_rows(),
        )
        logger.info(f"Service function {self.service_func_name} is created. Deployment completed successfully!")
        return service_spec_string, service_function_sql
