from collections.abc import Mapping, Sequence
from typing import Any, Literal, Optional
from uuid import UUID

from faker import Faker
from typing_extensions import deprecated

from ai.backend.common.arch import DEFAULT_IMAGE_ARCH
from ai.backend.common.typed_validators import SESSION_NAME_MAX_LENGTH
from ai.backend.common.types import RuntimeVariant

from ..exceptions import BackendClientError
from ..output.fields import service_fields
from ..output.types import FieldSpec, PaginatedResult
from ..pagination import fetch_paginated_result
from ..request import Request
from ..session import api_session
from ..utils import dedent as _d
from .base import BaseFunction, api_function

__all__ = ("Service",)

_default_fields: Sequence[FieldSpec] = (
    service_fields["endpoint_id"],
    service_fields["name"],
    service_fields["image_object"],
    service_fields["replicas"],
    service_fields["routings"],
    service_fields["url"],
    service_fields["session_owner"],
    service_fields["open_to_public"],
)


class Service(BaseFunction):
    id: UUID

    @api_function
    @classmethod
    @deprecated("Use paginated_list() instead of this method unless you set the name filter.")
    async def list(cls, name: Optional[str] = None) -> list[dict[str, Any]]:
        params = {}
        if name:
            params["name"] = name
        rqst = Request("GET", "/services", params=params)
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    @classmethod
    async def paginated_list(
        cls,
        *,
        fields: Sequence[FieldSpec] = _default_fields,
        page_offset: int = 0,
        page_size: int = 20,
        filter: Optional[str] = None,
        order: Optional[str] = None,
    ) -> PaginatedResult[dict[str, Any]]:
        return await fetch_paginated_result(
            "endpoint_list",
            {
                "filter": (filter, "String"),
                "order": (order, "String"),
            },
            fields,
            page_size=page_size,
            page_offset=page_offset,
        )

    @api_function
    @classmethod
    async def detail(
        cls,
        service_id: str,
        fields: Sequence[FieldSpec] = _default_fields,
    ) -> Sequence[dict[str, Any]]:
        query = _d("""
            query($endpoint_id: UUID!) {
                endpoint(endpoint_id: $endpoint_id) { $fields }
            }
        """)
        query = query.replace("$fields", " ".join(f.field_ref for f in fields))
        variables = {"endpoint_id": service_id}
        data = await api_session.get().Admin._query(query, variables)
        return data["endpoint"]

    @api_function
    @classmethod
    async def create(
        cls,
        image: str,
        model_id_or_name: str,
        initial_session_count: int,
        *,
        resources: Mapping[str, str | int],
        resource_opts: Mapping[str, str | int],
        domain_name: str,
        group_name: str,
        scaling_group: str,
        extra_mounts: Optional[Sequence[str]] = None,
        extra_mount_map: Optional[Mapping[str, str]] = None,
        extra_mount_options: Optional[Mapping[str, Mapping[str, str]]] = None,
        service_name: Optional[str] = None,
        model_version: Optional[str] = None,
        dependencies: Optional[Sequence[str]] = None,
        model_mount_destination: Optional[str] = None,
        envs: Optional[Mapping[str, str]] = None,
        startup_command: Optional[str] = None,
        cluster_size: int = 1,
        cluster_mode: Literal["single-node", "multi-node"] = "single-node",
        bootstrap_script: Optional[str] = None,
        tag: Optional[str] = None,
        architecture: Optional[str] = DEFAULT_IMAGE_ARCH,
        owner_access_key: Optional[str] = None,
        model_definition_path: Optional[str] = None,
        expose_to_public: bool = False,
        runtime_variant: Optional[RuntimeVariant] = None,
    ) -> dict[str, Any]:
        """
        Creates an inference service.

        :param image: The image name and tag for the infernence session.
            Example: ``python:3.6-ubuntu``.
            Check out the full list of available images in your server using (TODO:
            new API).
        :param service_name: A client-side (user-defined) identifier to distinguish the session among currently
            running sessions.
            It may be used to seamlessly reuse the session already created.
        :param initial_session_count: Number of sessions to be started along with
            service initiation.
        :param mounts: The ID of model type vFolder which contains model files required
            to start inference session.
        :param model_mount_destination: Path inside the container to mount model vFolder,
            defaults to /models
        :param envs: The environment variables which always bypasses the jail policy.
        :param resources: The resource specification. (TODO: details)
        :param cluster_size: The number of containers in this compute session.
            Must be at least 1.
        :param cluster_mode: Set the clustering mode whether to use distributed
            nodes or a single node to spawn multiple containers for the new session.
        :param tag: An optional string to annotate extra information.
        :param owner: An optional access key that owns the created session. (Only
            available to administrators)
        :param model_definition_path: Relative path to model definition file. Defaults to `model-definition.yaml`.
        :param expose_to_public: Visibility of API Endpoint which serves inference workload.
            If set to true, no authentication will be required to access the endpoint.
        :param runtime_variant: The runtime variant to use for the service.

        :returns: The :class:`ComputeSession` instance.
        """
        extra_mounts = extra_mounts or []
        extra_mount_map = extra_mount_map or {}
        extra_mount_options = extra_mount_options or {}

        if service_name is None:
            faker = Faker()
            service_name = f"bai-serve-{faker.user_name()}"[:SESSION_NAME_MAX_LENGTH]

        extra_mount_body = {}
        if extra_mounts:
            vfolder_id_to_name: dict[UUID, str] = {}
            vfolder_name_to_id: dict[str, UUID] = {}

            rqst = Request("GET", "/folders")
            async with rqst.fetch() as resp:
                body = await resp.json()
                for folder_info in body:
                    vfolder_id_to_name[UUID(folder_info["id"])] = folder_info["name"]
                    vfolder_name_to_id[folder_info["name"]] = UUID(folder_info["id"])

            for mount in extra_mounts:
                try:
                    vfolder_id = UUID(mount)
                    if vfolder_id not in vfolder_id_to_name:
                        raise BackendClientError(f"VFolder (id: {vfolder_id}) not found")
                except ValueError:
                    if mount not in vfolder_name_to_id:
                        raise BackendClientError(f"VFolder (name: {mount}) not found")
                    vfolder_id = vfolder_name_to_id[mount]
                extra_mount_body[str(vfolder_id)] = {
                    "mount_destination": extra_mount_map.get(mount),
                }
                if mount_type := extra_mount_options.get(mount, {}).get("type"):
                    extra_mount_body[str(vfolder_id)]["type"] = mount_type
        model_config = {
            "model": model_id_or_name,
            "model_mount_destination": model_mount_destination,
            "extra_mounts": extra_mount_body,
            "environ": envs,
            "scaling_group": scaling_group,
            "resources": resources,
            "resource_opts": resource_opts,
            "model_definition_path": model_definition_path,
        }
        if model_version:
            model_config["model_version"] = model_version
        model_config = {k: v for k, v in model_config.items() if v is not None}
        rqst = Request("POST", "/services")
        rqst_body = {
            "name": service_name,
            "replicas": initial_session_count,
            "image": image,
            "arch": architecture,
            "group": group_name,
            "domain": domain_name,
            "cluster_size": cluster_size,
            "cluster_mode": cluster_mode,
            "tag": tag,
            "startup_command": startup_command,
            "bootstrap_script": bootstrap_script,
            "owner_access_key": owner_access_key,
            "open_to_public": expose_to_public,
            "runtime_variant": runtime_variant,
            "config": model_config,
        }
        rqst_body = {k: v for k, v in rqst_body.items() if v is not None}
        rqst.set_json(rqst_body)
        async with rqst.fetch() as resp:
            body = await resp.json()
            return {
                "endpoint_id": body["endpoint_id"],
                "name": service_name,
                "replicas": initial_session_count,
                "active_route_count": 0,
                "service_endpoint": None,
                "is_public": expose_to_public,
            }

    @api_function
    @classmethod
    async def try_start(
        cls,
        image: str,
        model_id_or_name: str,
        *,
        resources: Mapping[str, str | int],
        resource_opts: Mapping[str, str | int],
        domain_name: str,
        group_name: str,
        scaling_group: str,
        service_name: Optional[str] = None,
        model_version: Optional[str] = None,
        dependencies: Optional[Sequence[str]] = None,
        model_mount_destination: Optional[str] = None,
        envs: Optional[Mapping[str, str]] = None,
        startup_command: Optional[str] = None,
        cluster_size: int = 1,
        cluster_mode: Literal["single-node", "multi-node"] = "single-node",
        bootstrap_script: Optional[str] = None,
        tag: Optional[str] = None,
        architecture: Optional[str] = DEFAULT_IMAGE_ARCH,
        owner_access_key: Optional[str] = None,
        expose_to_public: bool = False,
    ) -> dict[str, Any]:
        """
        Tries to start an inference session and terminates immediately.

        :param image: The image name and tag for the infernence session.
            Example: ``python:3.6-ubuntu``.
            Check out the full list of available images in your server using (TODO:
            new API).
        :param service_name: A client-side (user-defined) identifier to distinguish the session among currently
            running sessions.
            It may be used to seamlessly reuse the session already created.
        :param mounts: The ID of model type vFolder which contains model files required
            to start inference session.
        :param model_mount_destination: Path inside the container to mount model vFolder,
            defaults to /models
        :param envs: The environment variables which always bypasses the jail policy.
        :param resources: The resource specification. (TODO: details)
        :param cluster_size: The number of containers in this compute session.
            Must be at least 1.
        :param cluster_mode: Set the clustering mode whether to use distributed
            nodes or a single node to spawn multiple containers for the new session.
        :param tag: An optional string to annotate extra information.
        :param owner: An optional access key that owns the created session. (Only
            available to administrators)
        :param expose_to_public: Visibility of API Endpoint which serves inference workload.
            If set to true, no authentication will be required to access the endpoint.

        :returns: The :class:`ComputeSession` instance.
        """
        if service_name is None:
            faker = Faker()
            service_name = f"bai-serve-{faker.user_name()}"[:SESSION_NAME_MAX_LENGTH]

        rqst = Request("POST", "/services/_/try")
        rqst.set_json({
            "name": service_name,
            "replicas": 1,
            "image": image,
            "arch": architecture,
            "group": group_name,
            "domain": domain_name,
            "cluster_size": cluster_size,
            "cluster_mode": cluster_mode,
            "tag": tag,
            "startup_command": startup_command,
            "bootstrap_script": bootstrap_script,
            "owner_access_key": owner_access_key,
            "open_to_public": expose_to_public,
            "config": {
                "model": model_id_or_name,
                "model_version": model_version,
                "model_mount_destination": model_mount_destination,
                "environ": envs,
                "scaling_group": scaling_group,
                "resources": resources,
                "resource_opts": resource_opts,
            },
        })
        async with rqst.fetch() as resp:
            body = await resp.json()
            return {
                "task_id": body["task_id"],
                "name": service_name,
            }

    def __init__(self, id: str | UUID) -> None:
        super().__init__()
        self.id = id if isinstance(id, UUID) else UUID(id)

    @api_function
    async def info(self) -> dict[str, Any]:
        rqst = Request("GET", f"/services/{self.id}")
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def delete(self) -> dict[str, Any]:
        rqst = Request("DELETE", f"/services/{self.id}")
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def sync(self) -> dict[str, Any]:
        rqst = Request("POST", f"/services/{self.id}/sync")
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def scale(self, to: int) -> dict[str, Any]:
        rqst = Request("POST", f"/services/{self.id}/scale")
        rqst.set_json({"to": to})
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def generate_api_token(self, duration: str) -> dict[str, Any]:
        rqst = Request("POST", f"/services/{self.id}/token")
        rqst.set_json({"duration": duration})
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def update_traffic_ratio(self, target_route_id: UUID, new_ratio: float) -> dict[str, Any]:
        rqst = Request("PUT", f"/services/{self.id}/routings/{target_route_id}")
        rqst.set_json({"traffic_ratio": new_ratio})
        async with rqst.fetch() as resp:
            return await resp.json()

    @api_function
    async def downscale_single_route(self, target_route_id: UUID) -> dict[str, Any]:
        rqst = Request("DELETE", f"/services/{self.id}/routings/{target_route_id}")
        async with rqst.fetch() as resp:
            return await resp.json()
