import base64
import json
import os
import time
import uuid
from typing import List, Dict, Optional, Any

from cachetools import TTLCache
from tencentcloud.common import credential
from tencentcloud.common.abstract_client import AbstractClient
from tencentcloud.common.common_client import CommonClient
from tencentcloud.common.profile.client_profile import ClientProfile

from mlflow.entities.model_registry import (
    RegisteredModel,
    ModelVersion,
    ModelVersionTag,
)
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_ALREADY_EXISTS
from mlflow.store.model_registry import (
    SEARCH_REGISTERED_MODEL_MAX_RESULTS_THRESHOLD,
    SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD,
)
from mlflow.store.model_registry.abstract_store import AbstractStore
from mlflow.utils.annotations import experimental
from mlflow.utils.search_utils import SearchModelUtils, SearchModelVersionUtils
from mlflow.store.entities.paged_list import PagedList
from mlflow.models.model import get_model_info
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException

def to_string(obj):
    if obj is None:
        return "None"
    if isinstance(obj, str):
        return obj
    if isinstance(obj, dict):
        return json.dumps(obj, indent=2)
    try:
        return json.dumps(obj, indent=2)
    except:
        return str(obj)


tencent_cloud_debug = os.getenv("TENCENTCLOUD_DEBUG", None)


def log_msg(msg):
    if tencent_cloud_debug:
        print(msg)


def _get_create_time_from_audit(audit):
    if audit is None or "CreatedAt" not in audit:
        return None
    dt_str = audit["CreatedAt"]
    return int(dt_str)


def _get_last_updated_time_from_audit(audit):
    if audit is None or "LastModifiedAt" not in audit:
        return None
    dt_str = audit["LastModifiedAt"]
    if dt_str is None:
        return _get_create_time_from_audit(audit)
    return int(dt_str)

def _get_description(resp):
    if resp is None or "Comment" not in resp:
        return None
    return resp["Comment"]


TCLAKE_MLFLOW_RUN_ID_KEY = "tclake.mlflow.run_id"
TCLAKE_MLFLOW_RUN_LINK_KEY = "tclake.mlflow.run_link"
TCLAKE_MLFLOW_DEPLOYMENT_JOB_ID_KEY = "tclake.mlflow.deployment_job_id"
TCLAKE_WEDATA_WORKSPACE_ID_KEY = "tclake.wedata.workspace_id"
TCLAKE_MLFLOW_DEPLOYMENT_STATUS_KEY = "tclake.mlflow.deployment_status"
TCLAKE_MLFLOW_MODEL_ID_KEY = "tclake.mlflow.model_id"
TCLAKE_MLFLOW_MODEL_SIGNATURE_KEY = "tclake.mlflow.model_version_signature"
TCLAKE_MLFLOW_MODEL_ARTIFACT_PATH_KEY = "tclake.mlflow.artifact_path"
TCCATALOG_IDENTIFIER_KEY = "tccatalog.identifier"

UN_DEPLOYMENT = "UnDeployment"

def _set_kv_to_properties(key, value, properties=None):
    if properties is None:
        properties = []
    if value is not None:
        properties.append({"Key": key, "Value": value})
    return properties


def _get_kv_from_properties(properties, key):
    if properties is None:
        return None
    for p in properties:
        if p["Key"] == key:
            return p["Value"]
    return None


def _set_run_id_to_properties(run_id, properties=None):
    return _set_kv_to_properties(TCLAKE_MLFLOW_RUN_ID_KEY, run_id, properties)


def _get_run_id_from_properties(properties):
    return _get_kv_from_properties(properties, TCLAKE_MLFLOW_RUN_ID_KEY)


def _set_run_link_to_properties(run_link, properties):
    return _set_kv_to_properties(TCLAKE_MLFLOW_RUN_LINK_KEY, run_link, properties)


def _get_run_link_from_properties(properties):
    return _get_kv_from_properties(properties, TCLAKE_MLFLOW_RUN_LINK_KEY)


def _get_model_id_from_properties(properties):
    return _get_kv_from_properties(properties, TCLAKE_MLFLOW_MODEL_ID_KEY)


def _add_deployment_job_id_to_properties(deployment_job_id, properties=None):
    properties = _set_kv_to_properties(TCLAKE_MLFLOW_DEPLOYMENT_JOB_ID_KEY, deployment_job_id, properties)
    return properties


def _add_deployment_status_to_properties(status, properties=None):
    if status is not None:
        properties = _set_kv_to_properties(TCLAKE_MLFLOW_DEPLOYMENT_STATUS_KEY, status, properties)
    return properties


def _add_model_id_to_properties(model_id, properties=None):
    if model_id is not None:
        properties = _set_kv_to_properties(TCLAKE_MLFLOW_MODEL_ID_KEY, model_id, properties)
    return properties


def _add_model_signature_to_properties(source, properties):
    log_msg("model source is {}".format(source))
    model_info = get_model_info(source)  
    signature = model_info.signature
    log_msg("model {} signature is {}".format(source, signature))
    if signature is None:
        raise MlflowException(f"Registered model signature is not found in source artifact location '{source}'") 
    sig_json = json.dumps(signature.to_dict()) 
    log_msg("model {} signature json is {}".format(source, sig_json))
    properties = _set_kv_to_properties(TCLAKE_MLFLOW_MODEL_SIGNATURE_KEY, sig_json, properties)
    return properties


def _add_model_artifact_path_to_properties(source, properties):
    log_msg("model source is {}".format(source))
    model_info = get_model_info(source)  
    path = model_info.artifact_path
    log_msg("model {} artifact_path is {}".format(source, path))
    properties = _set_kv_to_properties(TCLAKE_MLFLOW_MODEL_ARTIFACT_PATH_KEY, path, properties)
    return properties


def _get_tags(asset_tag_votes: List[Dict[str, str]]):
    tags: List[ModelVersionTag] = []
    for tag in asset_tag_votes:
        tags.append(ModelVersionTag(tag["LabelName"], tag["LabelValue"]))
    return tags


def _get_model_version(version):
    return str(version)


def _set_model_version(version):
    return int(version)


def _make_model(resp, tags: List[Dict[str, str]]):
    audit = resp["Audit"]
    return RegisteredModel(
        name=resp["Name"],
        creation_timestamp=_get_create_time_from_audit(audit),
        last_updated_timestamp=_get_last_updated_time_from_audit(audit),
        description=_get_description(resp),
        tags=_get_tags(tags)
    )


def _get_model_version_name(entity):
    return "{}.{}.{}".format(
        entity["CatalogName"], entity["SchemaName"], entity["ModelName"]
    )


def _make_model_version(entity, name):
    properties = entity.get("Properties", None)
    audit = entity.get("Audit", None)
    aliases = entity.get("Aliases", [])
    asset_tag_votes = entity.get("Tags", [])
    return ModelVersion(
        name=name,
        version=_get_model_version(entity["Version"]),
        creation_timestamp=_get_create_time_from_audit(audit),
        last_updated_timestamp=_get_last_updated_time_from_audit(audit),
        description=_get_description(entity),
        source=entity["Uri"],
        run_id=_get_run_id_from_properties(properties),
        tags=_get_tags(asset_tag_votes),
        run_link=_get_run_link_from_properties(properties),
        aliases=aliases,
        model_id=_get_model_id_from_properties(properties),
        status="READY",
    )


def get_tencent_cloud_headers():
    header_json = os.getenv("TENCENTCLOUD_HEADER_JSON", None)
    if header_json is None:
        return None
    return json.loads(header_json)


def get_tencent_cloud_client_profile():
    endpoint = os.getenv("TENCENTCLOUD_ENDPOINT", None)
    if endpoint is None:
        return None
    client_profile = ClientProfile()
    client_profile.httpProfile.endpoint = endpoint
    return client_profile


def get_tencent_cloud_tccatalog_client_profile():
    endpoint = os.getenv("TENCENTCLOUD_TCCATALOG_ENDPOINT", None)
    if endpoint is None:
        return None
    client_profile = ClientProfile()
    client_profile.httpProfile.endpoint = endpoint
    return client_profile


def _parse_page_token(page_token):
    decoded_token = base64.b64decode(page_token)
    parsed_token = json.loads(decoded_token)
    return parsed_token


def _create_page_token(offset, search_id):
    return base64.b64encode(
        json.dumps({"offset": offset, "search_id": search_id}).encode("utf-8")
    )


def _get_asset_tag(asset_tags: List[Dict[str, str]], tag: ModelVersionTag) -> Optional[List[Dict[str, str]]]:
    tags: List[Dict[str, str]] = []
    if len(asset_tags) == 0:
        return [{"LabelName": tag.key, "LabelValue": tag.value}]

    for item in asset_tags:
        if item["LabelName"] == tag.key and item["LabelValue"] == tag.value:
            return None
        if item["LabelName"] == tag.key:
            tags.append({"LabelName": tag.key, "LabelValue": tag.value})
            continue
        tags.append({"LabelName": item["LabelName"], "LabelValue": item["LabelValue"]})
    return tags


class WedataClient(AbstractClient):
    _apiVersion = '2025-10-10'
    _endpoint = 'wedata.tencentcloudapi.com'
    _service = 'wedata'


class TcCatalogClient(AbstractClient):
    _apiVersion = '2024-10-24'
    _endpoint = 'tccatalog.tencentcloudapi.com'
    _service = 'tccatalog'


@experimental
class TCLakeStore(AbstractStore):
    """
    Client for an Open Source Unity Catalog Server accessed via REST API calls.
    """

    def __init__(self, store_uri=None, tracking_uri=None):
        super().__init__(store_uri, tracking_uri)
        log_msg(
            "initializing tencent tclake client {} {}".format(store_uri, tracking_uri)
        )
        sid = os.getenv("KERNEL_WEDATA_CLOUD_SDK_SECRET_ID", "")
        if len(sid) == 0:
            raise MlflowException("KERNEL_WEDATA_CLOUD_SDK_SECRET_ID is not set")
        sk = os.getenv("KERNEL_WEDATA_CLOUD_SDK_SECRET_KEY", "")
        if len(sk) == 0:
            raise MlflowException("KERNEL_WEDATA_CLOUD_SDK_SECRET_KEY is not set")

        token = os.getenv("KERNEL_WEDATA_CLOUD_SDK_SECRET_TOKEN", None)

        self.workspace_id = os.getenv("WEDATA_WORKSPACE_ID", "")
        log_msg(str.format("wedata workspace id: {}", self.workspace_id))

        self.sub_account_uin = os.getenv("KERNEL_LOGIN_UIN", "")
        self.owner_uin = os.getenv("QCLOUD_UIN", "")
        
        client_profile = get_tencent_cloud_client_profile()
        tccatalog_client_profile = get_tencent_cloud_tccatalog_client_profile()
        cred = credential.Credential(sid, sk, token)
        parts = store_uri.split(":")
        if len(parts) < 2:
            raise MlflowException("set store_uri tclake:{region}")
        region = parts[1]
        self.client = WedataClient(
            cred, region, client_profile
        )
        self.tccatalog_client = TcCatalogClient(
            cred, region, tccatalog_client_profile
        )
        self.headers = get_tencent_cloud_headers()
        self.default_catalog_name = os.getenv(
            "TENCENTCLOUD_DEFAULT_CATALOG_NAME", "default"
        )
        self.default_schema_name = os.getenv(
            "TENCENTCLOUD_DEFAULT_SCHEMA_NAME", "default"
        )
        cache_size = int(os.getenv("TCLAKE_CACHE_SIZE", "100"))
        cache_ttl = int(os.getenv("TCLAKE_CACHE_TTL_SECS", "300"))
        self.cache = TTLCache(maxsize=cache_size, ttl=cache_ttl)
        log_msg(
            "initialized tencent tclake client successfully {} {} {} {} {}".format(
                region, client_profile, self.headers, cache_size, cache_ttl
            )
        )

    def _split_model_name(self, name):
        parts = name.split(".")
        if len(parts) == 1:
            return [self.default_catalog_name, self.default_schema_name, parts[0]]
        if len(parts) == 2:
            return [self.default_catalog_name, parts[0], parts[1]]
        if len(parts) == 3:
            return parts
        raise MlflowException(
            "invalid model name: {}, must be catalog.schema.model".format(name)
        )

    def _create_labels(self, tags: List[Dict[str, str]]):
        already_created_tags = self._list_labels_by_name([tag["LabelName"] for tag in tags])

        create_label_req_body = {
            "WorkspaceId": self.workspace_id,
            "Labels": [],
        }
        create_label_value_req_body: Dict[str, Dict[str, Any]] = {}
        for tag in tags:
            # 标签已存在
            if tag["LabelName"] in already_created_tags and tag["LabelValue"] in already_created_tags[tag["LabelName"]]["Values"]:
                continue
            # 标签key已存在, 标签value不存在
            if tag["LabelName"] in already_created_tags:
                if tag["LabelName"] not in create_label_value_req_body:
                    # 后端服务提供者要求把标签id转换为int
                    create_label_value_req_body[tag["LabelName"]] = {"Values": [], "LabelId": int(already_created_tags[tag["LabelName"]]["Id"])}
                create_label_value_req_body[tag["LabelName"]]["Values"].append({"Value": tag["LabelValue"]})
                continue
            create_label_req_body["Labels"].append({
                "Name": tag["LabelName"],
                "Control": 2,
                "Type": 3,
                "Values": [{
                    "Value": tag["LabelValue"],
                }],
            })

        # 创建标签
        if len(create_label_req_body["Labels"]) != 0:
            resp = self._call("CreateLabels", create_label_req_body)
            log_msg(str.format("create labels response: {}", resp))

        # 创建标签值
        for body in create_label_value_req_body.values():
            resp = self._call("CreateLabelValues", body)
            log_msg(str.format("create label values response: {}", resp))

    def _list_labels_by_name(self, label_names: List[str]) -> Dict[str, Any]:
        req_body = {
            "WorkspaceId": self.workspace_id,
            "LabelNames": label_names,
            "Types": [3] * len(label_names),
            "Page": {
                "PageSize": len(label_names),
                "PageNumber": 1,
            }
        }
        resp = self._call("ListLabels", req_body)
        labels: Dict[str, Any] = {
            label["Name"]: {
                "Values": [v["Value"] for v in label["Values"]],
                "Id": label["Id"]
            }
            for label in resp["Labels"]
        }
        return labels

    def _list_tag_by_asset(self, property_type: str, property_id: str) -> List[Dict[str, str]]:
        req_body = {
            "QueryItems": [{
                "PropertyType": property_type,
                "PropertyId": property_id,
            }]
        }
        resp = self._call("ListTagsByAsset", req_body)
        tags: List[Dict[str, str]] = []
        for result in resp.get("Results", []):
            tags.extend(result.get("Tags", []))
        return tags

    def _batch_vote_asset_tag(self, property_type: str, property_id: str, tags: List[Dict[str, str]]):
        """批量更新资源标签, **覆盖式更新**"""
        req_body = {
            "Votes": [{
                "PropertyType": property_type,
                "PropertyId": property_id,
                "Creator": self.sub_account_uin,
                "OwnerAccount": self.owner_uin,
                "Tags": tags
            }]
        }
        self._call("BatchVoteAssetTag", req_body)


    def _call(self, action, req):
        log_msg("req: {}\n{}".format(action, json.dumps(req, indent=2)))
        body = self.client.call(action, req, headers=self.headers)
        body_obj = json.loads(body)
        log_msg("body: {}\n{}".format(action, json.dumps(body_obj, indent=2)))
        resp = body_obj["Response"]["Data"]
        return resp

    def _tccatalog_call(self, action, req):
        log_msg("req: {}\n{}".format(action, json.dumps(req, indent=2)))
        body = self.tccatalog_client.call(action, req, headers=self.headers)
        body_obj = json.loads(body)
        log_msg("body: {}\n{}".format(action, json.dumps(body_obj, indent=2)))
        resp = body_obj["Response"]
        return resp

    def _get_model_version_numbers(self, name):
        # TODO: 改用wedata ListAllModelVersionNumbers接口
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            # "WorkspaceId": self.workspace_id,
        }
        # resp = self._call("ListAllModelVersionNumbers", req_body)
        resp = self._tccatalog_call("DescribeModelVersionNumbers", req_body)
        return resp["Versions"]

    def _list_model_versions(self, name):
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "MaxResults": "1000",
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("ListModelVersions", req_body)
        model_versions = resp["Items"]
        while len(model_versions) < int(resp["TotalCount"]) and resp.get("NextPageToken", ""):
            time.sleep(0.05)
            req_body["PageToken"] = resp["NextPageToken"]
            resp = self._call("ListModelVersions", req_body)
            model_versions.extend(resp["Items"])
        return model_versions


    def _get_model_version_by_alias(self, name, alias):
        model_versions = self._list_model_versions(name)

        model_version = None
        for mv in model_versions:
            if "Aliases" not in mv:
                continue
            if alias in mv["Aliases"]:
                model_version = mv
                break
        if model_version is None:
            return None
        return model_version

    def _get_latest_model_version(self, name):
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        versions = self._get_model_version_numbers(name)
        log_msg("versions: {}".format(versions))
        latest_version = versions[-1]
        log_msg("latest_version: {}".format(latest_version))
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(latest_version),
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("GetModelVersion", req_body)
        model_version = resp["ModelVersion"]
        return model_version

    def create_registered_model(self, name, tags=None, description=None, deployment_job_id=None):
        log_msg("create_registered_model {} {} {}".format(name, tags, description))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        properties = []
        _add_deployment_job_id_to_properties(deployment_job_id, properties)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "Comment": description if description else "",
            "Properties": properties,
            "WorkspaceId": self.workspace_id,
        }
        try:
            resp = self._call("CreateModel", req_body)
        except Exception as e:
            if isinstance(e, TencentCloudSDKException):
                if e.code == "FailedOperation.MetalakeAlreadyExistsError":
                    raise MlflowException(
                        f"Registered Model (name={name}) already exists.",
                        RESOURCE_ALREADY_EXISTS,
                    )
            raise

        if not tags:
            return _make_model(resp["Model"], [])

        model_tags = [{"LabelName": tag.key, "LabelValue": tag.value} for tag in tags]
        try:
            self._create_labels(model_tags)
            self._batch_vote_asset_tag("MODEL", resp["Model"]["AssetGuid"], model_tags)
        except Exception as e:
            log_msg("failed to add model tag: {}".format(e))
            return _make_model(resp["Model"], [])

        return _make_model(resp["Model"], model_tags)
    
    def update_registered_model(self, name, description, deployment_job_id=None):
        log_msg("update_registered_model {} {} {}".format(name, description, deployment_job_id))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "NewComment": description if description else "",
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("UpdateModelComment", req_body)
        try:
            tags = self._list_tag_by_asset("MODEL", resp["Model"]["AssetGuid"])
        except Exception as e:
            log_msg("failed to list model tags: {}".format(e))
            tags = []
        return _make_model(resp["Model"], tags)

    def rename_registered_model(self, name, new_name):
        log_msg("rename_register_model {} {}".format(name, new_name))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "NewName": new_name,
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("UpdateModelName", req_body)
        try:
            tags = self._list_tag_by_asset("MODEL", resp["Model"]["AssetGuid"])
        except Exception as e:
            log_msg("failed to list model tags: {}".format(e))
            tags = []
        return _make_model(resp["Model"], tags)

    def delete_registered_model(self, name):
        log_msg("delete_register_model {}".format(name))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("DeleteModel", req_body)
        if not resp["Result"]:
            raise MlflowException("Failed to delete model {}".format(name))

    def _fetch_all_models(self):
        # TODO: 改用wedata SearchModels接口
        req_body = {"Offset": 0, "Limit": 200, "SnapshotBased": True, "SnapshotId": ""}
        resp = self._tccatalog_call("SearchModels", req_body)
        total = resp["TotalCount"]
        model_list = resp["Models"]
        while len(model_list) < total:
            time.sleep(0.05)
            req_body["Offset"] += req_body["Limit"]
            req_body["SnapshotId"] = resp["SnapshotId"]
            resp = self._tccatalog_call("SearchModels", req_body)
            model_list.extend(resp["Models"])
            if len(resp["Models"]) < req_body["Limit"]:
                break
        return [_make_model(model, []) for model in model_list]

    def search_registered_models(
        self, filter_string=None, max_results=None, order_by=None, page_token=None
    ):
        log_msg(
            "search_registered_models {} {} {} {}".format(
                filter_string, max_results, order_by, page_token
            )
        )
        if not isinstance(max_results, int) or max_results < 1:
            raise MlflowException(
                "Invalid value for max_results. It must be a positive integer,"
                f" but got {max_results}",
                INVALID_PARAMETER_VALUE,
            )

        if max_results > SEARCH_REGISTERED_MODEL_MAX_RESULTS_THRESHOLD:
            raise MlflowException(
                "Invalid value for request parameter max_results. It must be at most "
                f"{SEARCH_REGISTERED_MODEL_MAX_RESULTS_THRESHOLD}, but got value {max_results}",
                INVALID_PARAMETER_VALUE,
            )
        if page_token is None:
            registered_models = self._fetch_all_models()
            filtered_rms = SearchModelUtils.filter(registered_models, filter_string)
            sorted_rms = SearchModelUtils.sort(filtered_rms, order_by)
            if len(sorted_rms) == 0:
                return PagedList([], None)
            search_id = "model_" + str(uuid.uuid4())
            page_token = _create_page_token(0, search_id)
            self.cache[search_id] = sorted_rms
            log_msg(
                "search_registered_models add cache {} {} {} {} {} {}".format(
                    filter_string,
                    max_results,
                    order_by,
                    page_token,
                    search_id,
                    len(sorted_rms),
                )
            )
        return self._get_page_list(page_token, max_results)

    def _get_registered_model(self, name):
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("GetModel", req)
        return resp["Model"]

    def get_registered_model(self, name):
        log_msg("get_registered_model {}".format(name))
        model = self._get_registered_model(name)
        tags = self._list_tag_by_asset("MODEL", model["AssetGuid"])
        return _make_model(model, tags)

    def get_latest_versions(self, name, stages=None):
        log_msg("get_latest_versions {} {}".format(name, stages))
        if stages is not None:
            raise NotImplementedError("Method not support stages")
        # 获取模型版本列表
        model_versions = self._list_model_versions(name)
        for mv in model_versions:
            # 获取模型版本标签
            mv["Tags"] = self._list_tag_by_asset("MODEL_VERSION", mv["AssetGuid"])
        return [_make_model_version(mv, name) for mv in model_versions]

    def set_registered_model_tag(self, name, tag):
        log_msg("set_registered_model_tag {} {}".format(name, tag))
        model = self._get_registered_model(name)
        asset_tags = self._list_tag_by_asset("MODEL", model["AssetGuid"])
        tags: List[Dict[str, str]] = _get_asset_tag(asset_tags, tag)
        if tags is None:
            return

        self._batch_vote_asset_tag("MODEL", model["AssetGuid"], tags)

    def delete_registered_model_tag(self, name, key):
        log_msg("delete_registered_model_tag {} {}".format(name, key))
        model = self._get_registered_model(name)
        already_exists_tags = self._list_tag_by_asset("MODEL", model["AssetGuid"])
        if key not in [tag["LabelName"] for tag in already_exists_tags]:
            return

        tags: List[Dict[str, str]] = [tag for tag in already_exists_tags if tag["LabelName"] != key]
        self._batch_vote_asset_tag("MODEL", model["AssetGuid"], tags)

    def create_model_version(
        self,
        name,
        source,
        run_id=None,
        tags=None,
        run_link=None,
        description=None,
        local_model_path=None,
        model_id=None,
    ):
        log_msg(
            "create_model_version {} {} {} {} {} {} {}".format(
                name, source, run_id, tags, run_link, description, local_model_path
            )
        )
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        properties = []
        _set_run_id_to_properties(run_id, properties)
        _set_run_link_to_properties(run_link, properties)
        _add_deployment_status_to_properties(UN_DEPLOYMENT, properties)
        _add_model_id_to_properties(model_id, properties)
        _add_model_signature_to_properties(source, properties)
        _add_model_artifact_path_to_properties(source, properties)

        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "Uri": source,
            "Comment": description if description else "",
            "Properties": properties,
            "WorkspaceId": self.workspace_id,
        }
        self._call("CreateModelVersion", req_body)

        model_version = self._get_latest_model_version(name)
        if not tags:
            return _make_model_version(model_version, name)

        model_tags = [{"LabelName": tag.key, "LabelValue": tag.value} for tag in tags]
        try:
            self._create_labels(model_tags)
            self._batch_vote_asset_tag("MODEL_VERSION", model_version["AssetGuid"], model_tags)
            model_version["Tags"] = model_tags
        except Exception as e:
            log_msg("failed to add model version tag: {}".format(e))

        return _make_model_version(model_version, name)

    def update_model_version(self, name, version, description):
        log_msg("update_model_version {} {} {}".format(name, version, description))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(version),
            "NewComment": description if description else "",
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("UpdateModelVersionComment", req_body)
        model_version = resp["ModelVersion"]
        model_version["Tags"] = self._list_tag_by_asset("MODEL_VERSION", model_version["AssetGuid"])
        return _make_model_version(model_version, name)

    def transition_model_version_stage(
        self, name, version, stage, archive_existing_versions
    ):
        raise NotImplementedError("Method not implemented")

    def delete_model_version(self, name, version):
        log_msg("delete_model_version {} {}".format(name, version))
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(version),
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("DeleteModelVersion", req_body)
        if not resp["Result"]:
            raise Exception(
                "Failed to delete model version {} {}".format(name, version)
            )

    def _get_model_version(self, name, version):
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req_body = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(version),
            "WorkspaceId": self.workspace_id,
        }
        resp = self._call("GetModelVersion", req_body)
        return resp["ModelVersion"]

    def get_model_version(self, name, version):
        log_msg("get_model_version {} {}".format(name, version))
        model_version = self._get_model_version(name, version)
        tags = self._list_tag_by_asset("MODEL_VERSION", model_version["AssetGuid"])
        model_version["Tags"] = tags

        return _make_model_version(model_version, name)

    def _fetch_all_model_versions(self):
        # 由于底层加入了workspace权限管控隔离以后，平铺查询当前workspace下的所有模型的能力难以支持, 因此这里依旧调用tccatalog接口
        # TODO: 改用wedata SearchModelVersions接口
        req_body = {"Offset": 0, "Limit": 200, "SnapshotBased": True, "SnapshotId": ""}
        resp = self._tccatalog_call("SearchModelVersions", req_body)
        total = resp["TotalCount"]
        model_version_list = resp["ModelVersions"]
        while len(model_version_list) < total:
            time.sleep(0.05)
            req_body["Offset"] += req_body["Limit"]
            req_body["SnapshotId"] = resp["SnapshotId"]
            resp = self._tccatalog_call("SearchModelVersions", req_body)
            model_version_list.extend(resp["ModelVersions"])
            if len(resp["ModelVersions"]) < req_body["Limit"]:
                break
        return [
            _make_model_version(model_version, _get_model_version_name(model_version))
            for model_version in model_version_list
        ]

    def search_model_versions(
        self, filter_string=None, max_results=None, order_by=None, page_token=None
    ):
        log_msg(
            "search_model_versions {} {} {} {}".format(
                filter_string, max_results, order_by, page_token
            )
        )
        if not isinstance(max_results, int) or max_results < 1:
            raise MlflowException(
                "Invalid value for max_results. It must be a positive integer,"
                f" but got {max_results}",
                INVALID_PARAMETER_VALUE,
            )

        if max_results > SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD:
            raise MlflowException(
                "Invalid value for request parameter max_results. It must be at most "
                f"{SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD}, but got value {max_results}",
                INVALID_PARAMETER_VALUE,
            )
        if page_token is None:
            model_versions = self._fetch_all_model_versions()
            filtered_mvs = SearchModelVersionUtils.filter(model_versions, filter_string)
            sorted_mvs = SearchModelVersionUtils.sort(
                filtered_mvs,
                order_by
                or ["last_updated_timestamp DESC", "name ASC", "version_number DESC"],
            )
            if len(sorted_mvs) == 0:
                return PagedList([], None)
            search_id = "model_version_" + str(uuid.uuid4())
            page_token = _create_page_token(0, search_id)
            self.cache[search_id] = sorted_mvs
            log_msg(
                "search_model_versions add cache {} {} {} {} {} {}".format(
                    filter_string,
                    max_results,
                    order_by,
                    page_token,
                    search_id,
                    len(sorted_mvs),
                )
            )

        return self._get_page_list(page_token, max_results)

    def _get_page_list(self, page_token, max_results):
        token_info = _parse_page_token(page_token)
        log_msg(
            "_get_page_list token_info {} {} {}".format(
                page_token, max_results, token_info
            )
        )
        sorted_mvs = self.cache.get(token_info["search_id"])
        if sorted_mvs is None:
            raise MlflowException(
                "Invalid page token: search id not found or expired",
                INVALID_PARAMETER_VALUE,
            )
        start_offset = token_info["offset"]
        final_offset = start_offset + max_results

        paginated_rms = sorted_mvs[start_offset : min(len(sorted_mvs), final_offset)]
        next_page_token = None
        if final_offset < len(sorted_mvs):
            next_page_token = _create_page_token(final_offset, token_info["search_id"])
        else:
            self.cache.pop(token_info["search_id"], None)
            log_msg(
                "pop cache {} {} {} {} {} {}".format(
                    token_info["search_id"],
                    start_offset,
                    final_offset,
                    len(sorted_mvs),
                    page_token,
                    next_page_token,
                )
            )
        return PagedList(paginated_rms, next_page_token)

    def set_model_version_tag(self, name, version, tag):
        log_msg("set_model_version_tag {} {} {}".format(name, version, tag))
        model_version = self._get_model_version(name, version)
        asset_tags = self._list_tag_by_asset("MODEL_VERSION", model_version["AssetGuid"])
        tags: List[Dict[str, str]] = _get_asset_tag(asset_tags, tag)
        if tags is None:
            return
        self._batch_vote_asset_tag("MODEL_VERSION", model_version["AssetGuid"], tags)

    def delete_model_version_tag(self, name, version, key):
        log_msg("delete_model_version_tag {} {} {}".format(name, version, key))
        model_version = self._get_model_version(name, version)
        already_exists_tags = self._list_tag_by_asset("MODEL_VERSION", model_version["AssetGuid"])
        if key not in [item["LabelName"] for item in already_exists_tags]:
            return

        tags: List[Dict[str, str]] = [tag for tag in already_exists_tags if tag["LabelName"] != key]
        self._batch_vote_asset_tag("MODEL_VERSION", model_version["AssetGuid"], tags)

    def set_registered_model_alias(self, name, alias, version):
        # 首先检查模型别名是否存在
        model_version = self._get_model_version_by_alias(name, alias)
        if model_version is not None:
            raise MlflowException(f"Alias {alias} already exists for model {name}:{model_version['Version']}, please delete it first")
        # 设置模型别名
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
        req = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(version),
            "AddedAliases": [alias],
            "WorkspaceId": self.workspace_id,
        }
        self._call("UpdateModelVersionAliases", req)

    def delete_registered_model_alias(self, name, alias):
        model_version = self.get_model_version_by_alias(name, alias)
        if model_version is None:
            raise MlflowException(f"Alias {alias} does not exist for model {name}")
        [catalog_name, schema_name, model_name] = self._split_model_name(name)
         
        req = {
            "CatalogName": catalog_name,
            "SchemaName": schema_name,
            "ModelName": model_name,
            "ModelVersion": _set_model_version(model_version.version),
            "RemovedAliases": [alias],
            "WorkspaceId": self.workspace_id,
        }
        self._call("UpdateModelVersionAliases", req)

    def get_model_version_by_alias(self, name, alias):
        log_msg("get_model_version_by_alias {} {}".format(name, alias))
        mv = self._get_model_version_by_alias(name, alias)
        log_msg("get_model_version_by_alias {} {} {}".format(name, alias, mv))
        if mv is None:
              raise MlflowException(
                    f"Registered model alias {alias} not found.", INVALID_PARAMETER_VALUE
                )
        mv["Tags"] = self._list_tag_by_asset("MODEL_VERSION", mv["AssetGuid"])
        return _make_model_version(mv, name)

    def get_model_version_download_uri(self, name, version):
        log_msg("get_model_version_download_uri {} {}".format(name, version))
        model_version = self.get_model_version(name, version)
        return model_version.source

