#!/usr/bin/env python3

# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import os
import re
from collections import namedtuple

try:
    from mlflow.exceptions import MlflowException
except ImportError:
    # mlflow가 설치되지 않은 환경에서도 import 가능하도록 처리
    class MlflowException(Exception):  # type: ignore[no-redef]
        """Fallback MlflowException for when mlflow is not installed."""

        pass


# 환경변수 키 목록 (단일 소스)
ENV_VAR_KEYS = {
    # 필수 환경변수
    "MLFLOW_S3_ENDPOINT_URL",
    "MLFLOW_TRACKING_URI",
    "AWS_ACCESS_KEY_ID",
    "AWS_SECRET_ACCESS_KEY",
    "TRITON_URL",
    "TRITON_MODEL_REPO",
    "AWS_BUCKET_NAME",
    # 선택적 환경변수
    "MODEL_NAME",
}


def check_env():
    """Check required environment variables."""
    # MODEL_NAME은 선택사항이므로 제외
    required_env_vars = ENV_VAR_KEYS - {"MODEL_NAME"}

    missing_vars = [var for var in required_env_vars if not os.environ.get(var)]

    # 누락된 환경 변수가 있다면 에러 발생
    if missing_vars:
        error_message = (
            f"다음 환경 변수들이 설정되지 않았습니다: {', '.join(missing_vars)}"
        )
        raise OSError(error_message)


def load_env(args):
    """Load environment variables from args or env."""
    # ENV_VAR_KEYS를 사용하여 환경변수 로드
    for env_var in ENV_VAR_KEYS:
        if env_var in args:
            os.environ[env_var] = args[env_var]


class _TritonConfig(dict):
    """
    Internal configuration class for Triton deployment.

    This class is not exposed in the public API (__all__).
    It handles environment variable validation, S3 URI parsing,
    and boto3 client initialization.
    """

    def __init__(self):
        super().__init__()
        self["triton_url"] = os.environ.get("TRITON_URL")
        self["triton_model_repo"] = os.environ.get("TRITON_MODEL_REPO")
        self["model_name"] = os.environ.get("MODEL_NAME")  # Optional

        if not self["triton_url"]:
            raise Exception("Check that environment variable TRITON_URL is set")

        if not self["triton_model_repo"]:
            raise Exception("Check that environment variable TRITON_MODEL_REPO is set")

        if self["triton_model_repo"].startswith("s3://"):
            self.s3_regex = re.compile(
                "s3://(http://|https://|)([0-9a-zA-Z\\-.]+):([0-9]+)/"
                "([0-9a-z.\\-]+)(((/[0-9a-zA-Z.\\-_]+)*)?)"
            )

            uri = self.parse_path(self["triton_model_repo"])
            protocol = "https://" if uri.protocol == "https://" else "http://"
            endpoint_url = None
            if uri.host_name != "" and uri.host_port != "":
                endpoint_url = f"{protocol}{uri.host_name}:{uri.host_port}"

            # boto3 handles AWS credentials
            try:
                import boto3

                self["s3"] = boto3.client("s3", endpoint_url=endpoint_url)
                self["s3_bucket"] = uri.bucket
                self["s3_prefix"] = uri.prefix
            except ImportError:
                raise ImportError(
                    "boto3 is required for S3 storage support. Install it with: pip install boto3"
                )
            self["triton_model_repo"] = (
                f"s3://{uri.bucket}/{uri.prefix}"
                if uri.prefix
                else f"s3://{uri.bucket}"
            )

    def parse_path(self, path):
        # Cleanup extra slashes
        clean_path = self.clean_path(path)

        # Get the bucket name and the object path. Return error if path is malformed
        match = self.s3_regex.fullmatch(clean_path)
        S3URI = namedtuple(
            "S3URI", ["protocol", "host_name", "host_port", "bucket", "prefix"]
        )
        if match:
            uri = S3URI(*match.group(1, 2, 3, 4, 5))
            if uri.prefix and uri.prefix[0] == "/":
                uri = uri._replace(prefix=uri.prefix[1:])
        else:
            bucket_start = clean_path.find("s3://") + len("s3://")
            bucket_end = clean_path.find("/", bucket_start)

            # If there isn't a slash, the address has only the bucket
            if bucket_end > bucket_start:
                bucket = clean_path[bucket_start:bucket_end]
                prefix = clean_path[bucket_end + 1 :]
            else:
                bucket = clean_path[bucket_start:]
                prefix = ""
            uri = S3URI("", "", "", bucket, prefix)

        if uri.bucket == "":
            raise MlflowException("No bucket name found in path: " + path)

        return uri

    def clean_path(self, s3_path):
        # Must handle paths with s3 prefix
        start = s3_path.find("s3://")
        path = ""
        if start != -1:
            path = s3_path[start + len("s3://") :]
            clean_path = "s3://"
        else:
            path = s3_path
            clean_path = ""

        # Must handle paths with https:// or http:// prefix
        https_start = path.find("https://")
        if https_start != -1:
            path = path[https_start + len("https://") :]
            clean_path += "https://"
        else:
            http_start = path.find("http://")
            if http_start != -1:
                path = path[http_start + len("http://") :]
                clean_path += "http://"

        # Remove trailing slashes
        rtrim_length = len(path.rstrip("/"))
        if rtrim_length == 0:
            raise MlflowException("Invalid bucket name: '" + path + "'")

        # Remove leading slashes
        ltrim_length = len(path) - len(path.lstrip("/"))
        if ltrim_length == len(path):
            raise MlflowException("Invalid bucket name: '" + path + "'")

        # Remove extra internal slashes
        true_path = path[ltrim_length : rtrim_length + 1]
        previous_slash = False
        for i in range(len(true_path)):
            if true_path[i] == "/":
                if not previous_slash:
                    clean_path += true_path[i]
                previous_slash = True
            else:
                clean_path += true_path[i]
                previous_slash = False

        return clean_path
