# Copyright 2022-2026 Xinference Holdings Pte. Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
import logging
import os
import platform
import random
import re
import string
import sys
import uuid
import weakref
from enum import Enum
from typing import Any, Dict, Generator, List, Optional, Tuple, Union

import orjson
from packaging.markers import Marker
from packaging.requirements import Requirement

from .._compat import BaseModel
from ..constants import (
    XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
    XINFERENCE_LOG_ARG_MAX_LENGTH,
)

logger = logging.getLogger(__name__)


class AbortRequestMessage(Enum):
    NOT_FOUND = 1
    DONE = 2
    NO_OP = 3


def truncate_log_arg(arg) -> str:
    s = str(arg)
    if len(s) > XINFERENCE_LOG_ARG_MAX_LENGTH:
        s = s[0:XINFERENCE_LOG_ARG_MAX_LENGTH] + "..."
    return s


def log_async(
    logger,
    level=logging.DEBUG,
    ignore_kwargs: Optional[List[str]] = None,
    log_exception=True,
):
    import time
    from functools import wraps
    from inspect import signature

    def decorator(func):
        func_name = func.__name__
        sig = signature(func)

        @wraps(func)
        async def wrapped(*args, **kwargs):
            request_id_str = kwargs.get("request_id")
            if not request_id_str:
                # sometimes `request_id` not in kwargs
                # we try to bind the arguments
                try:
                    bound_args = sig.bind_partial(*args, **kwargs)
                    arguments = bound_args.arguments
                except TypeError:
                    arguments = {}
                request_id_str = arguments.get("request_id", "")
            if not request_id_str:
                request_id_str = uuid.uuid1()
                if func_name == "text_to_image":
                    kwargs["request_id"] = request_id_str
            request_id_str = f"[request {request_id_str}]"
            formatted_args = ",".join(map(truncate_log_arg, args))
            formatted_kwargs = ",".join(
                [
                    "%s=%s" % (k, truncate_log_arg(v))
                    for k, v in kwargs.items()
                    if ignore_kwargs is None or k not in ignore_kwargs
                ]
            )
            logger.log(
                level,
                f"{request_id_str} Enter {func_name}, args: {formatted_args}, kwargs: {formatted_kwargs}",
            )
            start = time.time()
            try:
                ret = await func(*args, **kwargs)
                logger.log(
                    level,
                    f"{request_id_str} Leave {func_name}, elapsed time: {int(time.time() - start)} s",
                )
                return ret
            except Exception as e:
                if log_exception:
                    logger.error(
                        f"{request_id_str} Leave {func_name}, error: {e}, elapsed time: {int(time.time() - start)} s",
                        exc_info=True,
                    )
                else:
                    logger.log(
                        level,
                        f"{request_id_str} Leave {func_name}, error: {e}, elapsed time: {int(time.time() - start)} s",
                    )
                raise

        return wrapped

    return decorator


def log_sync(logger, level=logging.DEBUG, log_exception=True):
    import time
    from functools import wraps

    def decorator(func):
        @wraps(func)
        def wrapped(*args, **kwargs):
            formatted_args = ",".join(map(truncate_log_arg, args))
            formatted_kwargs = ",".join(
                map(lambda x: "%s=%s" % (x[0], truncate_log_arg(x[1])), kwargs.items())
            )
            logger.log(
                level,
                f"Enter {func.__name__}, args: {formatted_args}, kwargs: {formatted_kwargs}",
            )
            start = time.time()
            try:
                ret = func(*args, **kwargs)
                logger.log(
                    level,
                    f"Leave {func.__name__}, elapsed time: {int(time.time() - start)} s",
                )
                return ret
            except Exception as e:
                if log_exception:
                    logger.error(
                        f"Leave {func.__name__}, error: {e}, elapsed time: {int(time.time() - start)} s",
                        exc_info=True,
                    )
                else:
                    logger.log(
                        level,
                        f"Leave {func.__name__}, error: {e}, elapsed time: {int(time.time() - start)} s",
                    )
                raise

        return wrapped

    return decorator


# Replica suffix must not collide with bare model names that end in
# "-<digits>" (llama-2, phi-2, gpt-2, ...). See #5198.
_REPLICA_MODEL_UID_SEP = "-rep"
_REPLICA_MODEL_UID_RE = re.compile(
    rf"^(?P<uid>.+){_REPLICA_MODEL_UID_SEP}(?P<rep>\d+)$"
)


def iter_replica_model_uid(model_uid: str, replica: int) -> Generator[str, None, None]:
    """
    Generates all the replica model uids.
    """
    replica = int(replica)
    for rep_id in range(replica):
        yield build_replica_model_uid(model_uid, rep_id)


def build_replica_model_uid(model_uid: str, rep_id: int) -> str:
    """
    Build a replica model uid.

    Uses a reserved ``-rep{n}`` suffix so bare model names that themselves
    end in ``-<digits>`` (e.g. ``llama-2``) are not ambiguous with a replica
    index.
    """
    return f"{model_uid}{_REPLICA_MODEL_UID_SEP}{int(rep_id)}"


def parse_replica_model_uid(replica_model_uid: str) -> Tuple[str, int]:
    """
    Parse replica model uid to model uid and rep id.

    Returns ``(model_uid, -1)`` when the input has no replica suffix.
    """
    match = _REPLICA_MODEL_UID_RE.match(replica_model_uid)
    if match:
        return match.group("uid"), int(match.group("rep"))
    # No reserved suffix: treat as bare model uid. Do NOT fall back to the
    # legacy "-{n}" split here — that mis-parses real model names like
    # "llama-2" as ("llama", 2). UIDs built by the old builder therefore no
    # longer parse as replicas; migration paths that may still see them
    # (e.g. worker recovery files written by an older version) must handle
    # the legacy format explicitly via parse_legacy_replica_model_uid.
    return replica_model_uid, -1


_LEGACY_REPLICA_MODEL_UID_RE = re.compile(r"^(?P<uid>.+)-(?P<rep>\d+)$")


def parse_legacy_replica_model_uid(
    replica_model_uid: str,
) -> Optional[Tuple[str, int]]:
    """
    Parse the legacy ``{model_uid}-{n}`` replica format used before the
    reserved ``-rep{n}`` suffix was introduced.

    Only intended for migration paths (e.g. worker recovery files written
    by an older version). The legacy format is ambiguous by design — a bare
    model name like ``llama-2`` also matches — so callers must verify that
    the returned base uid actually refers to a known model before trusting
    the result. Returns ``None`` when the input does not match the legacy
    pattern.
    """
    match = _LEGACY_REPLICA_MODEL_UID_RE.match(replica_model_uid)
    if match:
        return match.group("uid"), int(match.group("rep"))
    return None


def is_valid_model_uid(model_uid: str) -> bool:
    model_uid = model_uid.strip()
    if not model_uid or len(model_uid) > 100:
        return False
    # Forbid user-supplied uids that look like a replica suffix so they
    # cannot collide with build_replica_model_uid output.
    if _REPLICA_MODEL_UID_RE.match(model_uid):
        return False
    return True


def gen_random_string(length: int) -> str:
    return "".join(random.sample(string.ascii_letters + string.digits, length))


def json_dumps(o):
    def _default(obj):
        if isinstance(obj, BaseModel):
            return obj.dict()
        raise TypeError

    return orjson.dumps(o, default=_default)


def purge_dir(d):
    if not os.path.exists(d) or not os.path.isdir(d):
        return
    for name in os.listdir(d):
        subdir = os.path.join(d, name)
        try:
            if (os.path.islink(subdir) and not os.path.exists(subdir)) or (
                len(os.listdir(subdir)) == 0
            ):
                logger.info("Remove empty directory: %s", subdir)
                os.rmdir(subdir)
        except Exception:
            pass


def parse_model_version(model_version: str, model_type: str) -> Tuple:
    results: List[str] = model_version.split("--")
    if model_type == "LLM":
        if len(results) != 4:
            raise ValueError(
                f"LLM model_version parses failed! model_version: {model_version}"
            )
        model_name = results[0]
        size = results[1]
        if not size.endswith("B"):
            raise ValueError(f"Cannot parse model_size_in_billions: {size}")
        size = size.rstrip("B")
        size_in_billions: Union[int, str] = size if "_" in size else int(size)
        model_format = results[2]
        quantization = results[3]
        return model_name, size_in_billions, model_format, quantization
    elif model_type == "embedding":
        assert len(results) > 0, "Embedding model_version parses failed!"
        return (results[0],)
    elif model_type == "rerank":
        assert len(results) > 0, "Rerank model_version parses failed!"
        return (results[0],)
    elif model_type == "image":
        assert 2 >= len(results) >= 1, "Image model_version parses failed!"
        return tuple(results)
    else:
        raise ValueError(f"Not supported model_type: {model_type}")


def merge_virtual_env_packages(
    base_packages: List[str], extra_packages: Optional[List[str]]
) -> List[str]:
    """
    Merge default virtualenv packages with user provided ones. Packages with the
    same name will be replaced by the user supplied version instead of appended.
    """

    def get_key(package: str) -> str:
        pkg_name = package.split(";", 1)[0].strip()
        if pkg_name.startswith("#system_") and pkg_name.endswith("#"):
            return pkg_name[len("#system_") : -1].lower()
        if pkg_name.startswith("#"):
            return pkg_name
        try:
            return Requirement(package).name.lower()
        except Exception:
            # fallback: strip version/url markers best effort
            for sep in ["@", "==", ">=", "<=", "~=", "!=", "[", " "]:
                if sep in pkg_name:
                    return pkg_name.split(sep, 1)[0].strip().lower()
            return pkg_name.lower()

    merged: List[str] = []
    index_map: Dict[str, int] = {}
    for pkg in base_packages:
        canonical_key = get_key(pkg)
        pkg_name, separator, marker = pkg.partition(";")
        key = (
            f"{canonical_key}; {marker.strip()}"
            if separator and pkg_name.strip().startswith("#system_")
            else canonical_key
        )
        if key in index_map:
            merged[index_map[key]] = pkg
        else:
            index_map[key] = len(merged)
            merged.append(pkg)

    if extra_packages:
        for pkg in extra_packages:
            key = get_key(pkg)
            matching_indexes = [
                index
                for index, base_pkg in enumerate(merged)
                if get_key(base_pkg) == key
            ]
            if matching_indexes:
                merged[matching_indexes[0]] = pkg
                for index in reversed(matching_indexes[1:]):
                    merged.pop(index)
            else:
                merged.append(pkg)

    return merged


def normalize_sglang_kernel_packages(
    packages: List[str],
) -> Tuple[List[str], bool]:
    """Drop the legacy kernel distribution when a recipe selects its successor.

    SGLang 0.5.11 renamed the ``sgl-kernel`` distribution to
    ``sglang-kernel`` while keeping the same ``sgl_kernel`` import package.
    Installing both distributions into one cached virtualenv leaves two owners
    for the same files.  Model recipes that explicitly select the new
    distribution therefore supersede the legacy engine default.

    Returns the normalized package list and whether the modern kernel was
    selected.  Callers use the latter to migrate cached environments and force
    dependency reconciliation for the renamed distribution.
    """

    def _requirement_name(package: str) -> Optional[str]:
        head = package.split(";", 1)[0].strip()
        try:
            return Requirement(head).name.lower().replace("_", "-")
        except Exception:
            # Bare wheel URLs are not valid ``Requirement`` instances.  The
            # legacy CUDA 13 engine dependency uses exactly this form.
            lowered = head.lower()
            if re.search(r"(?:^|/)sgl[_-]kernel-", lowered):
                return "sgl-kernel"
            return None

    package_names = [_requirement_name(package) for package in packages]
    modern_kernel_selected = "sglang-kernel" in package_names
    if not modern_kernel_selected:
        return packages, False

    normalized = [
        package
        for package, package_name in zip(packages, package_names)
        if package_name != "sgl-kernel"
    ]
    return normalized, True


def build_subpool_envs_for_virtual_env(
    envs: Optional[Dict[str, str]],
    enable_virtual_env: Optional[bool],
    virtual_env_manager: Any,
    model_engine: Optional[str] = None,
) -> Dict[str, str]:
    subpool_envs = {} if envs is None else envs.copy()
    if bool(enable_virtual_env) and virtual_env_manager is not None:
        from .virtual_env_manager import resolve_virtualenv_python_path

        venv_python = resolve_virtualenv_python_path(virtual_env_manager)
        if not venv_python:
            return subpool_envs
        venv_bin = os.path.dirname(venv_python)
        venv_path = os.path.dirname(venv_bin)
        current_path = subpool_envs.get("PATH") or os.environ.get("PATH", "")
        subpool_envs["PATH"] = os.pathsep.join([venv_bin, current_path or ""])
        subpool_envs["VIRTUAL_ENV"] = venv_path
        subpool_envs.setdefault(
            "FLASHINFER_NINJA_PATH", os.path.join(venv_bin, "ninja")
        )
        if sys.platform == "linux" and (model_engine or "").lower() in (
            "sglang",
            "vllm",
        ):
            import sysconfig

            # PyTorch CUDA wheels load NVIDIA runtime libraries from Python
            # packages.  The actor subprocess is exec'd directly instead of
            # through virtualenv activation, so make both the child and
            # inherited parent locations visible to the dynamic linker.  CUDA
            # 13 packages use either the legacy per-library layout or the
            # consolidated nvidia/cu13/lib layout, depending on the wheel.
            child_site_packages = virtual_env_manager.get_lib_path()
            parent_site_packages = sysconfig.get_path("purelib")
            nvidia_lib_paths = []
            for site_packages in (child_site_packages, parent_site_packages):
                for directory in ("cusparselt", "cu13"):
                    path = os.path.join(site_packages, "nvidia", directory, "lib")
                    if path not in nvidia_lib_paths:
                        nvidia_lib_paths.append(path)
            current_ld_library_path = subpool_envs.get(
                "LD_LIBRARY_PATH"
            ) or os.environ.get("LD_LIBRARY_PATH", "")
            if current_ld_library_path:
                nvidia_lib_paths.append(current_ld_library_path)
            subpool_envs["LD_LIBRARY_PATH"] = os.pathsep.join(nvidia_lib_paths)
    return subpool_envs


def apply_engine_virtualenv_settings(
    settings: Any,
    model_engine: Optional[str],
) -> None:
    """
    Apply engine-specific virtualenv settings (extra indexes, strategies).
    """
    if not model_engine:
        return
    from .virtual_env_manager import (
        get_engine_virtualenv_extra_index_urls,
        get_engine_virtualenv_index_strategy,
    )

    engine_extra_index_urls = get_engine_virtualenv_extra_index_urls(model_engine)
    engine_index_strategy = get_engine_virtualenv_index_strategy(model_engine)
    if settings.extra_index_url is None and engine_extra_index_urls:
        settings.extra_index_url = engine_extra_index_urls
    if settings.index_strategy is None and engine_index_strategy:
        settings.index_strategy = engine_index_strategy


def filter_virtualenv_packages_by_markers(
    packages: List[str],
    model_engine: Optional[str],
    cuda_version: Optional[str],
    target_sys_platform: Optional[str] = None,
) -> List[str]:
    """
    Filter virtualenv packages using custom markers and system placeholders.

    This keeps #system_*# entries intact while evaluating engine/cuda/platform
    markers for other packages.
    """
    if not packages:
        return []
    effective_sys_platform = target_sys_platform or sys.platform
    system_markers = {
        "#system_torch#",
        "#system_torchaudio#",
        "#system_torchvision#",
        "#system_torchcodec#",
        "#system_numpy#",
        "#system_pandas#",
    }

    def _marker_allows(marker: str) -> bool:
        marker = marker.strip().lower()
        if "#engine#" in marker or "#model_engine#" in marker:
            if not model_engine:
                return False
            engine_value = model_engine.lower()
            if (
                f'#engine# == "{engine_value}"' not in marker
                and f'#model_engine# == "{engine_value}"' not in marker
                and f"#model_engine# == '{engine_value}'" not in marker
            ):
                return False

        if op_versions := re.findall(
            r"\bcuda_version\b\s*([<>=!~]+)\s*['\"]([^'\"]+)['\"]", marker
        ):
            if not cuda_version:
                return False
            if not all(
                check_marker(op_version, "cuda_version", cuda_version)
                for op_version in op_versions
            ):
                return False
        for operator, expected in re.findall(
            r"\bsys_platform\b\s*(==|!=)\s*['\"]([^'\"]+)['\"]", marker
        ):
            if operator == "==" and effective_sys_platform != expected:
                return False
            if operator == "!=" and effective_sys_platform == expected:
                return False
        if (
            'platform_machine == "x86_64"' in marker
            and platform.machine().lower() != "x86_64"
        ):
            return False
        if (
            'platform_machine == "aarch64"' in marker
            and platform.machine().lower() != "aarch64"
        ):
            return False
        return True

    def _has_custom_marker(marker: str) -> bool:
        marker = marker.lower()
        return (
            "#engine#" in marker
            or "#model_engine#" in marker
            or "cuda_version" in marker
        )

    resolved_packages: List[str] = []
    for pkg in packages:
        if ";" in pkg:
            req, marker = pkg.split(";", 1)
            req = req.strip()
            marker = marker.strip()
            if req in system_markers:
                if _has_custom_marker(marker):
                    if _marker_allows(marker):
                        resolved_packages.append(req)
                    continue
                resolved_packages.append(pkg)
                continue
            if _has_custom_marker(marker):
                if _marker_allows(marker):
                    resolved_packages.append(req)
                continue
        resolved_packages.append(pkg)
    return resolved_packages


def rewrite_direct_url_packages_for_index(packages: List[str]) -> List[str]:
    """
    Rewrite direct wheel-URL requirements to ``name==version`` specs.

    When a private index is explicitly configured (e.g. an offline mirror
    inherited from pip.conf via ``inherit_pip_config``), direct wheel URLs
    would bypass the index and hit the public internet, which fails in
    air-gapped deployments. The private index carries these wheels, so
    resolving them as ``name==version`` (keeping any local version segment
    such as ``+cu130``) is equivalent and stays fully offline.

    Non-wheel URLs (git+ sources, sdists) and unparsable entries are
    returned unchanged.
    """
    import posixpath
    from urllib.parse import unquote, urlparse

    rewritten: List[str] = []
    for pkg in packages:
        spec, sep, marker = pkg.partition(";")
        candidate = spec.strip()
        # Accept both "https://…/x.whl" and PEP 508 "name @ https://…/x.whl".
        if candidate.startswith(("http://", "https://")):
            url = candidate
        elif "@" in candidate:
            url = candidate.partition("@")[2].strip()
        else:
            rewritten.append(pkg)
            continue
        if not url.startswith(("http://", "https://")):
            rewritten.append(pkg)
            continue
        filename = unquote(posixpath.basename(urlparse(url).path))
        if not filename.endswith(".whl"):
            rewritten.append(pkg)
            continue
        # PEP 427: {name}-{version}(-{build})?-{python}-{abi}-{platform}.whl,
        # dashes inside name/version are escaped, so the first two fields
        # are exact.
        parts = filename[: -len(".whl")].split("-")
        if len(parts) < 5:
            rewritten.append(pkg)
            continue
        name, version = parts[0], parts[1]
        if not name or not version or not version[0].isdigit():
            rewritten.append(pkg)
            continue
        new_spec = f"{name}=={version}"
        rewritten.append(f"{new_spec} ; {marker.strip()}" if sep else new_spec)
    return rewritten


def find_direct_reference_packages(packages: List[str]) -> List[str]:
    """Return requirements that still bypass package indexes.

    Wheel URLs supported by :func:`rewrite_direct_url_packages_for_index`
    disappear before this helper is called. Remaining HTTP(S), ``git+`` and
    PEP 508 direct references cannot be satisfied reliably by an offline
    simple index, so callers can fail before an installer attempts egress.
    """

    direct_references: List[str] = []
    for pkg in packages:
        candidate = pkg.partition(";")[0].strip()
        if candidate.startswith(("http://", "https://", "git+")):
            direct_references.append(pkg)
            continue
        if "@" in candidate:
            target = candidate.partition("@")[2].strip()
            if target.startswith(("http://", "https://", "git+")):
                direct_references.append(pkg)
    return direct_references


def assign_replica_gpu(
    _replica_model_uid: str, replica: int, gpu_idx: Optional[Union[int, List[int]]]
) -> Optional[List[int]]:
    model_uid, rep_id = parse_replica_model_uid(_replica_model_uid)
    rep_id, replica = int(rep_id), int(replica)
    if isinstance(gpu_idx, int):
        gpu_idx = [gpu_idx]
    if isinstance(gpu_idx, list) and gpu_idx:
        if len(gpu_idx) % replica != 0:
            raise ValueError(
                "gpu_idx length must be a multiple of replica when specifying GPUs."
            )
        num_gpus = len(gpu_idx)
        gpus_per_replica = num_gpus // replica
        start = rep_id * gpus_per_replica
        end = start + gpus_per_replica
        return gpu_idx[start:end]
    return gpu_idx


class CancelMixin:
    _CANCEL_TASK_NAME = "abort_block"

    def __init__(self):
        self._running_tasks: weakref.WeakValueDictionary[  # type: ignore
            str, asyncio.Task
        ] = weakref.WeakValueDictionary()

    def _add_running_task(self, request_id: Optional[str]):
        """Add current asyncio task to the running task.
        :param request_id: The corresponding request id.
        """
        if request_id is None:
            return
        running_task = self._running_tasks.get(request_id)
        if running_task is not None:
            if running_task.get_name() == self._CANCEL_TASK_NAME:
                raise Exception(f"The request has been aborted: {request_id}")
            raise Exception(f"Duplicate request id: {request_id}")
        current_task = asyncio.current_task()
        assert current_task is not None
        self._running_tasks[request_id] = current_task

    def _cancel_running_task(
        self,
        request_id: Optional[str],
        block_duration: int = XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
    ):
        """Cancel the running asyncio task.
        :param request_id: The request id to cancel.
        :param block_duration: The duration seconds to ensure the request can't be executed.
        """
        if request_id is None:
            return
        running_task = self._running_tasks.pop(request_id, None)
        if running_task is not None:
            running_task.cancel()

        async def block_task():
            """This task is for blocking the request for a duration."""
            try:
                await asyncio.sleep(block_duration)
                logger.info("Abort block end for request: %s", request_id)
            except asyncio.CancelledError:
                logger.info("Abort block is cancelled for request: %s", request_id)

        if block_duration > 0:
            logger.info("Abort block start for request: %s", request_id)
            self._running_tasks[request_id] = asyncio.create_task(
                block_task(), name=self._CANCEL_TASK_NAME
            )


def check_marker(
    op_version: Tuple[str, str], package_name: str, package_version: str
) -> bool:
    try:
        op, version = op_version
        marker_str = f'python_full_version {op} "{version}"'
        env = {"python_full_version": package_version}
        return Marker(marker_str).evaluate(env)
    except Exception as e:
        logger.warning(
            f"Failed to evaluate {package_name} version marker {op_version} against {package_version}: {e}"
        )
        return False
