"""
通用工具函数
"""
import os
from collections import Counter
from datetime import datetime, timezone
from functools import wraps
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse

from mlflow.exceptions import RestException
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
from mlflow.utils import databricks_utils

from wedata.common.constants import constants
from wedata.common.constants.constants import MODEL_DATA_PATH_ROOT
from wedata.common.log import get_logger
from pyspark.sql import SparkSession


def validate_table_name(name: str):
    """
    验证特征表名规范，仅支持单表名，不能包含点（如<catalog>.<schema>.<table>）

    参数:
        name: 要验证的表名

    异常:
        ValueError: 如果表名包含点或不符合规范
    """
    if not name or not isinstance(name, str):
        raise ValueError("Table name must be a non-empty string")
    if name.count('.') > 0:
        raise ValueError("Feature table name only supports single table name, cannot contain dots (e.g. <catalog>.<schema>.<table>)")
    if not name[0].isalpha():
        raise ValueError("Table name must start with a letter")
    if not all(c.isalnum() or c == '_' for c in name):
        raise ValueError("Table name can only contain letters, numbers and underscores")


def build_full_table_name(
    table_name: str,
    database_name: Optional[str] = None,
    catalog_name: Optional[str] = None,
    spark_client: Optional[Any] = None
) -> str:
    """
    构建完整的表名，格式化为`<catalog>.<database>.<table>`或`<database>.<table>`形式。

    Args:
        table_name: 输入的表名（可以是简化的表名或完整表名）。
        database_name: 数据库名
        catalog_name: catalog名称（可选）。如果提供，将构建三级表名；如果不提供且spark_client存在，
                     将使用当前catalog；否则只构建二级表名。
        spark_client: SparkClient实例（可选）。用于获取当前catalog。

    Returns:
        完整表名（`<catalog>.<database>.<table>`或`<database>.<table>`）。
    """

    feature_store_database_name = os.environ.get("WEDATA_DEFAULT_FEATURE_STORE_DATABASE")
    logger = get_logger()
    if database_name:
        feature_store_database_name = database_name

    if not feature_store_database_name:
        logger.error("The current user has not configured a default feature database. "
                     "Please contact the manager account to configure it.")
        raise RuntimeError("Feature store is not configured! Please contact the main account to configure it.")

    logger.debug("feature database:{}".format(feature_store_database_name))

    # 构建基础的 database.table 名称
    base_table_name = f"{feature_store_database_name}.{table_name}"

    # 如果提供了 catalog_name，构建三级表名
    if catalog_name:
        return f"{catalog_name}.{base_table_name}"

    # 如果没有提供 catalog_name 但提供了 spark_client，使用当前 catalog
    if spark_client is not None:
        try:
            current_catalog = spark_client.get_current_catalog()
            if current_catalog:
                return f"{current_catalog}.{base_table_name}"
        except Exception as e:
            logger.warning(f"Failed to get current catalog: {e}. Using 2-level table name.")

    # 默认返回二级表名（向后兼容）
    return base_table_name


def enable_if(condition):
    """
    A decorator that conditionally enables a function based on a condition.
    If the condition is not truthy, calling the function raises a NotImplementedError.

    :param condition: A callable that returns a truthy or falsy value.
    """

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not condition():
                raise NotImplementedError
            return func(*args, **kwargs)

        return wrapper

    return decorator


def is_empty(target: str):
    return target is None or len(target.strip()) == 0


class _NoDbutilsError(Exception):
    pass


def _get_dbutils():
    try:
        import IPython

        ip_shell = IPython.get_ipython()
        if ip_shell is None:
            raise _NoDbutilsError
        return ip_shell.ns_table["user_global"]["dbutils"]
    except ImportError:
        raise _NoDbutilsError
    except KeyError:
        raise _NoDbutilsError


def utc_timestamp_ms_from_iso_datetime_string(date_string: str) -> int:
    dt = datetime.fromisoformat(date_string)
    utc_dt = dt.replace(tzinfo=timezone.utc)
    return int(1000 * utc_dt.timestamp())


def pip_depependency_pinned_major_version(pip_package_name, major_version):
    """
    Generate a pip dependency string that is pinned to a major version, for example: "databricks-feature-lookup==0.*"
    """
    return f"{pip_package_name}=={major_version}.*"

def pip_depependency_pinned_version(pip_package_name, version):
    """
    Generate a pip dependency string that is pinned to a major version, for example: "databricks-feature-lookup==0.*"
    """
    return f"{pip_package_name}=={version}"


def add_mlflow_pip_depependency(conda_env, pip_package_name):
    """
    Add a new pip dependency to the conda environment taken from the raw MLflow model.
    """
    if pip_package_name is None or len(pip_package_name) == 0:
        raise ValueError(
            "Unexpected input: missing or empty pip_package_name parameter"
        )

    found_pip_dependency = False
    if conda_env is not None:
        for dep in conda_env["dependencies"]:
            if isinstance(dep, dict) and "pip" in dep:
                found_pip_dependency = True
                pip_deps = dep["pip"]
                if pip_package_name not in pip_deps:
                    pip_deps.append(pip_package_name)
        if "dependencies" in conda_env and not found_pip_dependency:
            raise ValueError(
                "Unexpected input: mlflow conda_env did not contain pip as a dependency"
            )


def download_model_artifacts(model_uri, dir):
    """
    Downloads model artifacts from model_uri to dir.
    """
    if not is_artifact_uri(model_uri):
        raise ValueError(
            f"Invalid model URI '{model_uri}'."
            f"Use ``models:/model_name>/<version_number>`` or "
            f"``runs:/<mlflow_run_id>/run-relative/path/to/model``."
        )

    try:
        repo = get_artifact_repository(model_uri)
    except RestException as e:
        raise ValueError(f"The model at '{model_uri}' does not exist.", e)

    artifact_path = os.path.join("artifacts", MODEL_DATA_PATH_ROOT)
    if len(repo.list_artifacts(artifact_path)) == 0:
        raise ValueError(
            f"No suitable model found at '{model_uri}'. Either no model exists in this "
            f"artifact location or an existing model was not packaged with Feature Store metadata. "
            f"Only models logged by FeatureStoreClient.log_model can be used in inference."
        )

    return repo.download_artifacts(artifact_path="", dst_path=dir)


def validate_params_non_empty(params: Dict[str, Any], expected_params: List[str]):
    """
    Validate that none of the expected parameters are empty.
    """
    for expected_param in expected_params:
        if expected_param not in params:
            raise ValueError(
                f'Internal error: expected parameter "{expected_param}" not found in params dictionary'
            )
        param_value = params[expected_param]
        if not param_value:
            raise ValueError(f'Parameter "{expected_param}" cannot be empty')


def get_workspace_url() -> Optional[str]:
    """
    Overrides the behavior of the mlflow.utils.databricks_utils.get_workspace_url().
    """
    workspace_url = databricks_utils.get_workspace_url()
    if workspace_url and not urlparse(workspace_url).scheme:
        workspace_url = "https://" + workspace_url
    return workspace_url


def is_artifact_uri(uri):
    """
    Checks the artifact URI is associated with a MLflow model or run.
    The actual URI can be a model URI, model URI + subdirectory, or model URI + path to artifact file.
    """
    return ModelsArtifactRepository.is_models_uri(
        uri
    ) or RunsArtifactRepository.is_runs_uri(uri)


def as_list(obj, default=None):
    if not obj:
        return default
    elif isinstance(obj, list):
        return obj
    else:
        return [obj]


def get_duplicates(elements: List[Any]) -> List[Any]:
    """
    Returns duplicate elements in the order they first appear.
    """
    element_counts = Counter(elements)
    duplicates = []
    for e in element_counts.keys():
        if element_counts[e] > 1:
            duplicates.append(e)
    return duplicates


def validate_strings_unique(strings: List[str], error_template: str):
    """
    Validates all strings are unique, otherwise raise ValueError with the error template and duplicates.
    Passes single-quoted, comma delimited duplicates to the error template.
    """
    duplicate_strings = get_duplicates(strings)
    if duplicate_strings:
        duplicates_formatted = ", ".join([f"'{s}'" for s in duplicate_strings])
        raise ValueError(error_template.format(duplicates_formatted))


def sanitize_identifier(identifier: str):
    """
    Sanitize and wrap an identifier with backquotes. For example, "a`b" becomes "`a``b`".
    Use this function to sanitize identifiers such as column names in SQL and PySpark.
    """
    return f"`{identifier.replace('`', '``')}`"


def sanitize_identifiers(identifiers: List[str]):
    """
    Sanitize and wrap the identifiers in a list with backquotes.
    """
    return [sanitize_identifier(i) for i in identifiers]


def sanitize_multi_level_name(multi_level_name: str):
    """
    Sanitize a multi-level name (such as an Unity Catalog table name) by sanitizing each segment
    and joining the results. For example, "ca+t.fo`o.ba$r" becomes "`ca+t`.`fo``o`.`ba$r`".
    """
    segments = multi_level_name.split(".")
    return ".".join(sanitize_identifiers(segments))


def unsanitize_identifier(identifier: str):
    """
    Unsanitize an identifier. Useful when we get a possibly sanitized identifier from Spark or
    somewhere else, but we need an unsanitized one.
    Note: This function does not check the correctness of the identifier passed in. e.g. `foo``
    is not a valid sanitized identifier. When given such invalid input, this function returns
    invalid output.
    """
    if len(identifier) >= 2 and identifier[0] == "`" and identifier[-1] == "`":
        return identifier[1:-1].replace("``", "`")
    else:
        return identifier


# strings containing \ or ' can break sql statements, so escape them.
def escape_sql_string(input_str: str) -> str:
    return input_str.replace("\\", "\\\\").replace("'", "\\'")


def get_unique_list_order(elements: List[Any]) -> List[Any]:
    """
    Returns unique elements in the order they first appear.
    """
    return list(dict.fromkeys(elements))


def validate_database(database_name):
    if database_name is None:
        database_name = os.environ.get(constants.WEDATA_DEFAULT_FEATURE_STORE_DATABASE)
    if database_name is None:
        raise ValueError("Database_name variable or default database is not set.")
    return True


def check_package_version(package_name, expected_version, op="=="):
    """
    检查指定包的版本是否满足预期版本要求。
    Args:
        package_name: 包名称
        expected_version: 预期版本要求，例如3.5.5
        op: 比较运算符，默认为 "=="
    Returns:
        (是否成功找到包，版本是否匹配，已安装版本)
    如果满足，返回 (True, True, installed_version)；否则返回 (True, False, installed_version)。
    如果指定包不存在，返回 (False, False, None)。
    """
    # 在脚本顶部添加
    from packaging import version
    import importlib.metadata
    try:
        installed_version = importlib.metadata.version(package_name)

        if not op:
            raise ValueError(f"Invalid op: {op}. need be in ['==', '>', '<', '>=', '<=', '!=', '~=']")
        # 支持版本范围检查（如 ">=2.0,<3.0"）
        # 使用 packaging.version 进行复杂版本`检查
        i = version.parse(installed_version)
        e = version.parse(expected_version)
        return True, eval(f"i{op}e"), installed_version

    except importlib.metadata.PackageNotFoundError:
        return False, False, None


def check_spark_table_exists(spark_client: SparkSession, full_table_name: str) -> bool:
    _, ok, _ = check_package_version("pyspark", "3.5.0", ">=")

    # 优先使用 SQL 查询，避免 catalog API 在某些环境下的兼容性问题
    # 例如：DataLakeCatalog 可能会将 database 名称误认为 catalog 名称
    try:
        split = full_table_name.split(".")
        if len(split) == 2:
            # database.table 格式
            query = f"SHOW TABLES IN {split[0]} LIKE '{split[1]}'"
        elif len(split) == 3:
            # catalog.database.table 格式
            query = f"SHOW TABLES IN {split[0]}.{split[1]} LIKE '{split[2]}'"
        else:
            # 只有 table 名称
            query = f"SHOW TABLES LIKE '{full_table_name}'"
        return spark_client.sql(query).count() > 0
    except Exception as e:
        # 如果 SQL 查询失败，尝试使用 catalog API（向后兼容）
        try:
            return spark_client.catalog.tableExists(full_table_name)
        except Exception:
            # 如果两种方法都失败，抛出原始异常
            raise e
