"""
API Handler for calling Vertex AI Model Garden Models

Most Vertex Model Garden Models are OpenAI compatible - so this handler calls `openai_like_chat_completions`

Usage:

response = litellm.completion(
    model="vertex_ai/openai/5464397967697903616",
    messages=[{"role": "user", "content": "Hello, how are you?"}],
)

Sent to this route when `model` is in the format `vertex_ai/openai/{MODEL_ID}`


Vertex Documentation for using the OpenAI /chat/completions endpoint: https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_deployment.ipynb
"""

from collections.abc import Callable
from typing import Final

import httpx

from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.utils import ModelResponse

from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..vertex_llm_base import VertexBase


def _vertex_model_garden_model_id_in_json_body(model: str) -> bool:
    """
    Vertex catalog / publisher models are addressed as publisher/model (e.g.
    xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body.

    Deployed Model Garden endpoints are typically a single segment (often numeric)
    and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field.
    """
    return "/" in model


def create_vertex_url(
    vertex_location: str,
    vertex_project: str,
    stream: bool | None,
    model: str,
    api_base: str | None = None,
) -> str:
    """Return the api base for vertex model garden (without /chat/completions)."""
    base_url: Final = get_vertex_base_url(vertex_location)
    if _vertex_model_garden_model_id_in_json_body(model):
        return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi"
    return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}"


class VertexAIModelGardenModels(VertexBase):
    def __init__(self) -> None:
        super().__init__()

    def completion(
        self,
        model: str,
        messages: list,
        model_response: ModelResponse,
        print_verbose: Callable,
        encoding,
        logging_obj,
        api_base: str | None,
        optional_params: dict,
        custom_prompt_dict: dict,
        headers: dict | None,
        timeout: float | httpx.Timeout,
        litellm_params: dict,
        vertex_project=None,
        vertex_location=None,
        vertex_credentials=None,
        logger_fn=None,
        acompletion: bool = False,
        client=None,
    ):
        """
        Handles calling Vertex AI Model Garden Models in OpenAI compatible format

        Sent to this route when `model` is in the format `vertex_ai/openai/{MODEL_ID}`
        """
        try:
            import vertexai

            from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler
        except Exception as e:
            raise VertexAIError(
                status_code=400,
                message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""",
            )

        if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")):
            raise VertexAIError(
                status_code=400,
                message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
            )
        try:
            model = get_vertex_base_model_name(model=model)

            access_token, project_id = self._ensure_access_token(
                credentials=vertex_credentials,
                project_id=vertex_project,
                custom_llm_provider="vertex_ai",
            )

            openai_like_chat_completions: Final = OpenAILikeChatHandler()

            ## CONSTRUCT API BASE
            # Skip _check_custom_proxy: its ":verb" URL construction corrupts a
            # user-supplied api_base (e.g. Vertex MG dedicated endpoint), and
            # OpenAILikeChatHandler already appends "/chat/completions".
            stream: Final[bool] = optional_params.get("stream", False) or False
            optional_params["stream"] = stream
            if api_base is None:
                api_base = create_vertex_url(
                    vertex_location=vertex_location or "us-central1",
                    vertex_project=vertex_project or project_id,
                    stream=stream,
                    model=model,
                )
            # Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route).
            # Single-segment endpoint ids: model is encoded in the URL path; body model stays empty.
            if not _vertex_model_garden_model_id_in_json_body(model):
                model = ""
            return openai_like_chat_completions.completion(
                model=model,
                messages=messages,
                api_base=api_base,
                api_key=access_token,
                custom_prompt_dict=custom_prompt_dict,
                model_response=model_response,
                print_verbose=print_verbose,
                logging_obj=logging_obj,
                optional_params=optional_params,
                acompletion=acompletion,
                litellm_params=litellm_params,
                logger_fn=logger_fn,
                client=client,
                timeout=timeout,
                encoding=encoding,
                custom_llm_provider="vertex_ai",
            )

        except Exception as e:
            raise VertexAIError(status_code=500, message=str(e))