#!/usr/bin/env python3
import asyncio
import dataclasses
import difflib
from concurrent import futures
import argparse
import signal
import sys
import os
import json
import time
import gc
import tempfile
from typing import List
from PIL import Image

import backend_pb2
import backend_pb2_grpc

import grpc
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
from grpc_auth import get_auth_interceptors

from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.sampling_params import SamplingParams
from vllm.utils import random_uuid
try:
    from vllm.tokenizers import get_tokenizer  # vLLM >= 0.22
except ImportError:
    from vllm.transformers_utils.tokenizer import get_tokenizer  # vLLM < 0.22
from vllm.multimodal.utils import fetch_image
from vllm.assets.video import VideoAsset
import base64
import io

# Version-compat imports — wrap in try/except for older vLLM versions
try:
    from vllm.tool_parsers import ToolParserManager
    HAS_TOOL_PARSERS = True
except ImportError:
    HAS_TOOL_PARSERS = False

try:
    from vllm.reasoning import ReasoningParserManager
    HAS_REASONING_PARSERS = True
except ImportError:
    HAS_REASONING_PARSERS = False

# vLLM >= 0.23 renamed GuidedDecodingParams -> StructuredOutputsParams and the
# SamplingParams field guided_decoding -> structured_outputs.
try:
    from vllm.sampling_params import StructuredOutputsParams
    HAS_GUIDED_DECODING = True
except ImportError:
    HAS_GUIDED_DECODING = False

_ONE_DAY_IN_SECONDS = 60 * 60 * 24

# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))

# Implement the BackendServicer class with the service methods
class BackendServicer(backend_pb2_grpc.BackendServicer):
    """
    A gRPC servicer that implements the Backend service defined in backend.proto.
    """
    def generate(self,prompt, max_new_tokens):
        """
        Generates text based on the given prompt and maximum number of new tokens.

        Args:
            prompt (str): The prompt to generate text from.
            max_new_tokens (int): The maximum number of new tokens to generate.

        Returns:
            str: The generated text.
        """
        self.generator.end_beam_search()

        # Tokenizing the input
        ids = self.generator.tokenizer.encode(prompt)

        self.generator.gen_begin_reuse(ids)
        initial_len = self.generator.sequence[0].shape[0]
        has_leading_space = False
        decoded_text = ''
        for i in range(max_new_tokens):
            token = self.generator.gen_single_token()
            if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):
                has_leading_space = True

            decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])
            if has_leading_space:
                decoded_text = ' ' + decoded_text

            if token.item() == self.generator.tokenizer.eos_token_id:
                break
        return decoded_text

    def _parse_options(self, options_list):
        """Parse Options[] key:value string list into a dict."""
        opts = {}
        for opt in options_list:
            if ":" not in opt:
                continue
            key, value = opt.split(":", 1)
            opts[key.strip()] = value.strip()
        return opts

    def _apply_engine_args(self, engine_args, engine_args_json):
        """Apply user-supplied engine_args (JSON object) onto an AsyncEngineArgs.

        Returns a new AsyncEngineArgs with the typed fields preserved and the
        user's overrides layered on top. Uses ``dataclasses.replace`` so vLLM's
        ``__post_init__`` re-runs and auto-converts dict-valued fields like
        ``compilation_config`` / ``attention_config`` into their dataclass form.
        ``speculative_config`` and ``kv_transfer_config`` are accepted as dicts
        directly (vLLM converts them at engine init).

        Unknown keys raise ValueError with the closest valid field as a hint.
        """
        if not engine_args_json:
            return engine_args
        try:
            extra = json.loads(engine_args_json)
        except json.JSONDecodeError as e:
            raise ValueError(f"engine_args is not valid JSON: {e}") from e
        if not isinstance(extra, dict):
            raise ValueError(
                f"engine_args must be a JSON object, got {type(extra).__name__}"
            )
        valid = {f.name for f in dataclasses.fields(type(engine_args))}
        for key in extra:
            if key not in valid:
                suggestion = difflib.get_close_matches(key, valid, n=1)
                hint = f" did you mean {suggestion[0]!r}?" if suggestion else ""
                raise ValueError(f"unknown engine_args key {key!r}.{hint}")
        return dataclasses.replace(engine_args, **extra)

    def _messages_to_dicts(self, messages):
        """Convert proto Messages to list of dicts suitable for apply_chat_template()."""
        result = []
        for msg in messages:
            d = {"role": msg.role, "content": msg.content or ""}
            if msg.name:
                d["name"] = msg.name
            if msg.tool_call_id:
                d["tool_call_id"] = msg.tool_call_id
            if msg.reasoning_content:
                d["reasoning_content"] = msg.reasoning_content
            if msg.tool_calls:
                try:
                    tool_calls = json.loads(msg.tool_calls)
                except json.JSONDecodeError:
                    pass
                else:
                    # OpenAI wire format carries function.arguments as a
                    # JSON-encoded string, but chat templates (e.g. Qwen3)
                    # iterate over it as a mapping. vLLM's own OpenAI server
                    # parses arguments before applying the template, so do
                    # the same here.
                    if isinstance(tool_calls, list):
                        for tc in tool_calls:
                            func = tc.get("function") if isinstance(tc, dict) else None
                            if isinstance(func, dict) and isinstance(func.get("arguments"), str):
                                try:
                                    func["arguments"] = json.loads(func["arguments"])
                                except json.JSONDecodeError:
                                    pass
                    d["tool_calls"] = tool_calls
            result.append(d)
        return result

    def Health(self, request, context):
        """
        Returns a health check message.

        Args:
            request: The health check request.
            context: The gRPC context.

        Returns:
            backend_pb2.Reply: The health check reply.
        """
        return backend_pb2.Reply(message=bytes("OK", 'utf-8'))

    async def LoadModel(self, request, context):
        """
        Loads a language model.

        Args:
            request: The load model request.
            context: The gRPC context.

        Returns:
            backend_pb2.Result: The load model result.
        """
        engine_args = AsyncEngineArgs(
            model=request.Model,
        )

        if request.Quantization != "":
            engine_args.quantization = request.Quantization
        if request.LoadFormat != "":
            engine_args.load_format = request.LoadFormat
        if request.GPUMemoryUtilization != 0:
            engine_args.gpu_memory_utilization = request.GPUMemoryUtilization
        if request.TrustRemoteCode:
            engine_args.trust_remote_code = request.TrustRemoteCode
        if request.EnforceEager:
            engine_args.enforce_eager = request.EnforceEager
        if request.TensorParallelSize:
            engine_args.tensor_parallel_size = request.TensorParallelSize
        if request.SwapSpace != 0:
            engine_args.swap_space = request.SwapSpace
        if request.MaxModelLen != 0:
            engine_args.max_model_len = request.MaxModelLen
        if request.DisableLogStatus:
            engine_args.disable_log_status = request.DisableLogStatus
        if request.DType != "":
            engine_args.dtype = request.DType
        if request.LimitImagePerPrompt != 0 or request.LimitVideoPerPrompt != 0 or request.LimitAudioPerPrompt != 0:
            # limit-mm-per-prompt defaults to 1 per modality, based on vLLM docs
            engine_args.limit_mm_per_prompt = {
                "image": max(request.LimitImagePerPrompt, 1),
                "video": max(request.LimitVideoPerPrompt, 1),
                "audio": max(request.LimitAudioPerPrompt, 1)
            }

        # engine_args from YAML overrides typed fields above so operators can
        # tune anything the AsyncEngineArgs dataclass exposes without waiting
        # on protobuf changes.
        try:
            engine_args = self._apply_engine_args(engine_args, request.EngineArgs)
        except ValueError as err:
            print(f"engine_args error: {err}", file=sys.stderr)
            return backend_pb2.Result(success=False, message=str(err))

        try:
            self.llm = AsyncLLMEngine.from_engine_args(engine_args)
        except Exception as err:
            print(f"Unexpected {err=}, {type(err)=}", file=sys.stderr)
            return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")

        try:
            # vLLM >= 0.14 removed get_model_config() on AsyncLLM; the tokenizer
            # is either already loaded on the engine or can be built from the
            # Model name directly.
            tokenizer = None
            if hasattr(self.llm, "get_tokenizer"):
                try:
                    tokenizer = await self.llm.get_tokenizer()
                except TypeError:
                    tokenizer = self.llm.get_tokenizer()
                except Exception:
                    tokenizer = None
            if tokenizer is None and hasattr(self.llm, "tokenizer"):
                tokenizer = self.llm.tokenizer
            if tokenizer is None:
                tokenizer = get_tokenizer(
                    request.Model,
                    trust_remote_code=bool(request.TrustRemoteCode),
                    truncation_side="left",
                )
            self.tokenizer = tokenizer
        except Exception as err:
            return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")

        # Parse options for parser selection
        opts = self._parse_options(request.Options)

        # Instantiate tool/reasoning parser classes (they'll be instantiated per-request with tokenizer)
        self.tool_parser_cls = None
        self.reasoning_parser_cls = None
        if HAS_TOOL_PARSERS and opts.get("tool_parser"):
            try:
                self.tool_parser_cls = ToolParserManager.get_tool_parser(opts["tool_parser"])
                print(f"Loaded tool_parser: {opts['tool_parser']}", file=sys.stderr)
            except Exception as e:
                print(f"Failed to load tool_parser {opts.get('tool_parser')}: {e}", file=sys.stderr)

        if HAS_REASONING_PARSERS and opts.get("reasoning_parser"):
            try:
                self.reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(opts["reasoning_parser"])
                print(f"Loaded reasoning_parser: {opts['reasoning_parser']}", file=sys.stderr)
            except Exception as e:
                print(f"Failed to load reasoning_parser {opts.get('reasoning_parser')}: {e}", file=sys.stderr)

        print("Model loaded successfully", file=sys.stderr)
        return backend_pb2.Result(message="Model loaded successfully", success=True)

    async def Predict(self, request, context):
        """
        Generates text based on the given prompt and sampling parameters.

        Args:
            request: The predict request.
            context: The gRPC context.

        Returns:
            backend_pb2.Reply: The predict result.
        """
        gen = self._predict(request, context, streaming=False)
        res = await gen.__anext__()
        return res

    def Embedding(self, request, context):
        """
        A gRPC method that calculates embeddings for a given sentence.

        Args:
            request: An EmbeddingRequest object that contains the request parameters.
            context: A grpc.ServicerContext object that provides information about the RPC.

        Returns:
            An EmbeddingResult object that contains the calculated embeddings.
        """
        print("Calculated embeddings for: " + request.Embeddings, file=sys.stderr)
        outputs = self.model.encode(request.Embeddings)
        # Check if we have one result at least
        if len(outputs) == 0:
            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
            context.set_details("No embeddings were calculated.")
            return backend_pb2.EmbeddingResult()
        return backend_pb2.EmbeddingResult(embeddings=outputs[0].outputs.embedding)

    async def PredictStream(self, request, context):
        """
        Generates text based on the given prompt and sampling parameters, and streams the results.

        Args:
            request: The predict stream request.
            context: The gRPC context.

        Returns:
            backend_pb2.Result: The predict stream result.
        """
        iterations = self._predict(request, context, streaming=True)
        try:
            async for iteration in iterations:
                yield iteration
        finally:
            await iterations.aclose()

    async def TokenizeString(self, request, context):
        if not hasattr(self, 'tokenizer') or self.tokenizer is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Model/tokenizer not loaded")
            return backend_pb2.TokenizationResponse()
        try:
            tokens = self.tokenizer.encode(request.Prompt)
            return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
        except Exception as e:
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(str(e))
            return backend_pb2.TokenizationResponse()

    async def Free(self, request, context):
        try:
            if hasattr(self, 'llm'):
                del self.llm
            if hasattr(self, 'tokenizer'):
                del self.tokenizer
            self.tool_parser_cls = None
            self.reasoning_parser_cls = None
            gc.collect()
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass
            return backend_pb2.Result(success=True, message="Model freed")
        except Exception as e:
            return backend_pb2.Result(success=False, message=str(e))

    async def Score(self, request, context):
        """
        Joint log-probability of each candidate continuation given the
        shared prompt. Used by routing-policy multi-label classification
        (read the distribution rather than asking the model to emit a
        single argmax label), reranking, and reward-model scoring.

        Implementation uses vLLM's `prompt_logprobs` to recover the
        per-token log P(token_i | tokens_<i) for the full concatenated
        sequence; the candidate's tokens are the suffix whose logprobs
        get summed. max_tokens=1 because vLLM requires at least one
        generated token; the generated token is discarded.
        """
        if not hasattr(self, 'llm') or self.llm is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Model not loaded")
            return backend_pb2.ScoreResponse()
        if not hasattr(self, 'tokenizer') or self.tokenizer is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Tokenizer not available")
            return backend_pb2.ScoreResponse()
        if len(request.candidates) == 0:
            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
            context.set_details("candidates must be non-empty")
            return backend_pb2.ScoreResponse()

        try:
            prompt = request.prompt or ""
            prompt_token_ids = self.tokenizer.encode(prompt)
            prompt_len = len(prompt_token_ids)
            results = []

            for candidate in request.candidates:
                # Tokenise the concatenated sequence. We can't naively
                # use len(prompt_tokens) + len(tokenizer.encode(candidate))
                # because BPE merges at the boundary may produce a
                # different tokenisation. Encoding the joined text and
                # walking the divergence point is the correct primitive.
                full_text = prompt + candidate
                full_token_ids = self.tokenizer.encode(full_text)

                divergence = prompt_len
                min_len = min(prompt_len, len(full_token_ids))
                for i in range(min_len):
                    if prompt_token_ids[i] != full_token_ids[i]:
                        divergence = i
                        break

                candidate_token_ids = full_token_ids[divergence:]
                num_candidate_tokens = len(candidate_token_ids)
                if num_candidate_tokens == 0:
                    results.append(backend_pb2.CandidateScore(
                        log_prob=0.0,
                        length_normalized_log_prob=0.0,
                        num_tokens=0,
                    ))
                    continue

                sampling = SamplingParams(
                    max_tokens=1,
                    temperature=0.0,
                    prompt_logprobs=1,
                    detokenize=False,
                )

                request_id = random_uuid()
                last_output = None
                outputs_iter = self.llm.generate(
                    {"prompt": full_text},
                    sampling_params=sampling,
                    request_id=request_id,
                )
                try:
                    async for out in outputs_iter:
                        last_output = out
                finally:
                    try:
                        await outputs_iter.aclose()
                    except Exception:
                        pass

                _pl = getattr(last_output, "prompt_logprobs", None) if last_output is not None else None
                # Some engines accept the prompt_logprobs request but return a
                # list of all-None entries instead of computing them (observed
                # with vllm-metal's MLX backend on macOS). Treat that as
                # unsupported rather than silently scoring every candidate as 0.
                if not _pl or all(e is None for e in _pl):
                    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
                    context.set_details("This backend did not return prompt_logprobs; scoring is unsupported on this engine (e.g. vllm-metal / MLX on macOS).")
                    return backend_pb2.ScoreResponse()

                prompt_logprobs = last_output.prompt_logprobs
                total = 0.0
                tokens_proto = []
                for offset, tok_id in enumerate(candidate_token_ids):
                    position = divergence + offset
                    if position >= len(prompt_logprobs) or prompt_logprobs[position] is None:
                        continue
                    entry = prompt_logprobs[position]
                    lp_obj = entry.get(tok_id)
                    if lp_obj is not None:
                        lp = lp_obj.logprob
                    else:
                        # Token not in top-K; vLLM's top-1 may miss it.
                        # Fall back to the lowest available logprob in the
                        # entry — a conservative lower-bound on the true
                        # log P, biased against this candidate.
                        lp = min(v.logprob for v in entry.values())
                    total += lp
                    if request.include_token_logprobs:
                        tokens_proto.append(backend_pb2.TokenLogProb(
                            token=self.tokenizer.decode([tok_id]),
                            log_prob=lp,
                        ))

                cs = backend_pb2.CandidateScore(
                    log_prob=total,
                    num_tokens=num_candidate_tokens,
                )
                if request.length_normalize and num_candidate_tokens > 0:
                    cs.length_normalized_log_prob = total / num_candidate_tokens
                if tokens_proto:
                    cs.tokens.extend(tokens_proto)
                results.append(cs)

            return backend_pb2.ScoreResponse(candidates=results)
        except Exception as e:
            print(f"Score error: {e}", file=sys.stderr)
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(str(e))
            return backend_pb2.ScoreResponse()

    async def _predict(self, request, context, streaming=False):
        # Build the sampling parameters
        # NOTE: this must stay in sync with the vllm backend
        request_to_sampling_params = {
            "N": "n",
            "PresencePenalty": "presence_penalty",
            "FrequencyPenalty": "frequency_penalty",
            "RepetitionPenalty": "repetition_penalty",
            "Temperature": "temperature",
            "TopP": "top_p",
            "TopK": "top_k",
            "MinP": "min_p",
            "Seed": "seed",
            "StopPrompts": "stop",
            "StopTokenIds": "stop_token_ids",
            "BadWords": "bad_words",
            "IncludeStopStrInOutput": "include_stop_str_in_output",
            "IgnoreEOS": "ignore_eos",
            "Tokens": "max_tokens",
            "MinTokens": "min_tokens",
            "Logprobs": "logprobs",
            "PromptLogprobs": "prompt_logprobs",
            "SkipSpecialTokens": "skip_special_tokens",
            "SpacesBetweenSpecialTokens": "spaces_between_special_tokens",
            "TruncatePromptTokens": "truncate_prompt_tokens",
        }

        sampling_params = SamplingParams(top_p=0.9, max_tokens=200)

        for request_field, param_field in request_to_sampling_params.items():
            if hasattr(request, request_field):
                value = getattr(request, request_field)
                if value not in (None, 0, [], False):
                    setattr(sampling_params, param_field, value)

        # Structured-output decoding: use Grammar field to pass JSON schema or BNF
        if HAS_GUIDED_DECODING and request.Grammar:
            try:
                json.loads(request.Grammar)  # valid JSON = JSON schema
                sampling_params.structured_outputs = StructuredOutputsParams(json=request.Grammar)
            except json.JSONDecodeError:
                sampling_params.structured_outputs = StructuredOutputsParams(grammar=request.Grammar)

        # Extract image paths and process images
        prompt = request.Prompt

        image_paths = request.Images
        image_data = [self.load_image(img_path) for img_path in image_paths]

        videos_path = request.Videos
        video_data = [self.load_video(video_path) for video_path in videos_path]

        # If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
        if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
            messages_dicts = self._messages_to_dicts(request.Messages)
            template_kwargs = {"tokenize": False, "add_generation_prompt": True}

            # Pass tools for tool calling
            if request.Tools:
                try:
                    template_kwargs["tools"] = json.loads(request.Tools)
                except json.JSONDecodeError:
                    pass

            # Enable thinking mode if requested
            if request.Metadata.get("enable_thinking", "").lower() == "true":
                template_kwargs["enable_thinking"] = True

            try:
                prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
            except TypeError:
                # Some tokenizers don't support tools/enable_thinking kwargs — retry without them
                prompt = self.tokenizer.apply_chat_template(
                    messages_dicts, tokenize=False, add_generation_prompt=True
                )

        # Generate text using the LLM engine
        request_id = random_uuid()
        print(f"Generating text with request_id: {request_id}", file=sys.stderr)
        multi_modal_data = {}
        if image_data:
            multi_modal_data["image"] = image_data
        if video_data:
            multi_modal_data["video"] = video_data
        outputs = self.llm.generate(
            {
            "prompt": prompt,
            "multi_modal_data": multi_modal_data if multi_modal_data else None,
            },
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Stream the results
        generated_text = ""
        generated_token_ids: list[int] = []
        last_output = None

        # Tool-parsing strategy decision (made once, before the loop):
        #
        # When a tool parser is active, the model's raw tool-call markup
        # (e.g. <tool_call>...) must not be streamed verbatim as delta.content
        # — clients would see the unparsed syntax. Two paths:
        #
        # (A) native streaming via parser.extract_tool_calls_streaming. All
        #     concrete tool parsers shipped with vLLM 0.23+ implement this
        #     (Granite4, Qwen3Coder, DeepSeekV31, Jamba, Ernie45, Hermes,
        #     llama3_json, mistral, …). The parser decides per-delta whether
        #     to emit content or suppress tool-call markup, and emits a
        #     structured DeltaMessage(tool_calls=[...]) when a call is ready.
        # (B) buffer fallback — used only when the parser surprisingly lacks
        #     the streaming method or it raises mid-stream. The post-loop
        #     extract_tool_calls assembles the final chat_delta. Same correctness
        #     guarantee as a non-streaming response, at the cost of a delayed
        #     final chunk.
        has_tool_parser = bool(self.tool_parser_cls and request.Tools)
        tp_instance = None
        tp_request = None
        native_streaming = False
        native_streaming_error = False
        if has_tool_parser:
            try:
                tools_for_parser = json.loads(request.Tools)
            except json.JSONDecodeError:
                tools_for_parser = []
            try:
                tp_instance = self.tool_parser_cls(self.tokenizer, tools=tools_for_parser)
            except TypeError:
                tp_instance = self.tool_parser_cls(self.tokenizer)
            # Build a minimal ChatCompletionRequest so the streaming method
            # sees the tools list. We do not need any other request fields —
            # parsers only read .tools (and sometimes .tool_choice, which we
            # leave at default).
            try:
                from vllm.entrypoints.openai.chat_completion.protocol import (
                    ChatCompletionRequest as _CCR,
                )
                tp_request = _CCR(
                    model="local",
                    messages=[{"role": "user", "content": ""}],
                    tools=tools_for_parser or None,
                )
            except Exception as e:
                print(f"Could not build ChatCompletionRequest for streaming parser: {e}",
                      file=sys.stderr)
                tp_request = None
            native_streaming = (
                tp_request is not None
                and hasattr(tp_instance, "extract_tool_calls_streaming")
            )

        try:
            async for request_output in outputs:
                iteration_text = request_output.outputs[0].text
                last_output = request_output

                if streaming:
                    delta_iteration_text = iteration_text.removeprefix(generated_text)
                    new_token_ids = list(request_output.outputs[0].token_ids)
                    delta_token_ids = new_token_ids[len(generated_token_ids):]

                    if not has_tool_parser:
                        # Plain streaming — unchanged from pre-tool-parser path.
                        yield backend_pb2.Reply(
                            message=bytes(delta_iteration_text, encoding='utf-8'),
                            chat_deltas=[backend_pb2.ChatDelta(content=delta_iteration_text)],
                        )
                    elif native_streaming and not native_streaming_error:
                        # (A) Native vLLM extract_tool_calls_streaming.
                        try:
                            msg = tp_instance.extract_tool_calls_streaming(
                                previous_text=generated_text,
                                current_text=iteration_text,
                                delta_text=delta_iteration_text,
                                previous_token_ids=generated_token_ids,
                                current_token_ids=new_token_ids,
                                delta_token_ids=delta_token_ids,
                                request=tp_request,
                            )
                        except Exception as e:
                            print(f"Streaming tool parser error (falling back to "
                                  f"buffer for the rest of the stream): {e}",
                                  file=sys.stderr)
                            native_streaming_error = True
                            msg = None
                        if msg is not None:
                            tc_protos = []
                            for tc in (msg.tool_calls or []):
                                fn = tc.function or None
                                tc_protos.append(backend_pb2.ToolCallDelta(
                                    index=tc.index,
                                    id=tc.id or "",
                                    name=(fn.name if fn and fn.name else "") or "",
                                    arguments=(fn.arguments if fn and fn.arguments else "") or "",
                                ))
                            cd_kwargs = {}
                            if msg.content:
                                cd_kwargs["content"] = msg.content
                            if msg.reasoning:
                                cd_kwargs["reasoning_content"] = msg.reasoning
                            if tc_protos:
                                cd_kwargs["tool_calls"] = tc_protos
                            if cd_kwargs:
                                yield backend_pb2.Reply(
                                    message=bytes(msg.content or "", encoding='utf-8'),
                                    chat_deltas=[backend_pb2.ChatDelta(**cd_kwargs)],
                                )
                    # (B) buffer fallback — emit nothing during the stream.
                    # The post-loop extract_tool_calls block builds the final chunk.

                # Keep track of text + token_ids generated
                generated_text = iteration_text
                generated_token_ids = list(request_output.outputs[0].token_ids)
        finally:
            await outputs.aclose()

        # Remove the image files from /tmp folder
        for img_path in image_paths:
            try:
                os.remove(img_path)
            except Exception as e:
                print(f"Error removing image file: {img_path}, {e}", file=sys.stderr)

        # Parse reasoning and tool calls from final text using vLLM's native parsers
        content = generated_text
        reasoning_content = ""
        tool_calls_proto = []

        if self.reasoning_parser_cls:
            try:
                rp = self.reasoning_parser_cls(self.tokenizer)
                r, c = rp.extract_reasoning(generated_text, request=None)
                reasoning_content = r or ""
                content = c if c is not None else generated_text
            except Exception as e:
                print(f"Reasoning parser error: {e}", file=sys.stderr)

        # When (A) native streaming ran cleanly, per-delta yields above already
        # delivered everything — do NOT extract again on the full text or we'd
        # duplicate content/tool_calls into the final chunk.
        # NOTE: `native_streaming` is a capability flag ("streaming parser is
        # available"), not a state flag ("streaming actually ran"). For
        # non-streaming requests it is still True but the per-delta loop was
        # never entered, so we MUST still run extract_tool_calls here. Hence
        # the explicit `streaming and …` guard on both branches.
        if has_tool_parser and not (streaming and native_streaming and not native_streaming_error):
            try:
                tp = tp_instance
                if tp is None:
                    # Defensive: tp_instance build failed earlier; reconstruct.
                    tools = json.loads(request.Tools)
                    try:
                        tp = self.tool_parser_cls(self.tokenizer, tools=tools)
                    except TypeError:
                        tp = self.tool_parser_cls(self.tokenizer)
                info = tp.extract_tool_calls(content, request=None)
                if info.tools_called:
                    content = info.content or ""
                    for i, tc in enumerate(info.tool_calls):
                        tool_calls_proto.append(backend_pb2.ToolCallDelta(
                            index=i,
                            id=tc.id,
                            name=tc.function.name,
                            arguments=tc.function.arguments,
                        ))
            except Exception as e:
                print(f"Tool parser error: {e}", file=sys.stderr)
        elif streaming and native_streaming and not native_streaming_error:
            # Per-delta path already emitted content + tool_calls; the final
            # chat_delta should carry only metadata (token counts, logprobs).
            content = ""

        # Extract token counts
        prompt_tokens = 0
        completion_tokens = 0
        if last_output is not None:
            try:
                prompt_tokens = len(last_output.prompt_token_ids or [])
            except Exception:
                pass
            try:
                completion_tokens = len(last_output.outputs[0].token_ids or [])
            except Exception:
                pass

        # Extract logprobs if requested
        logprobs_bytes = b""
        if last_output is not None and request.Logprobs > 0:
            try:
                lp = last_output.outputs[0].logprobs
                if lp:
                    logprobs_data = {"content": []}
                    for token_lp_dict in lp:
                        if token_lp_dict:
                            first_tok_id, first_lp = next(iter(token_lp_dict.items()))
                            logprobs_data["content"].append({
                                "token": getattr(first_lp, "decoded_token", str(first_tok_id)),
                                "logprob": first_lp.logprob,
                            })
                    logprobs_bytes = json.dumps(logprobs_data).encode("utf-8")
            except Exception as e:
                print(f"Logprobs extraction error: {e}", file=sys.stderr)

        chat_delta = backend_pb2.ChatDelta(
            content=content,
            reasoning_content=reasoning_content,
            tool_calls=tool_calls_proto,
        )

        if streaming:
            # Final chunk with structured data.
            #
            # If we used the buffer fallback (has_tool_parser=True AND native
            # streaming did NOT run cleanly) and the parser found no tool call,
            # flush the buffered content as ONE content delta — and clear the
            # final chat_delta's content so the metadata chunk does not repeat
            # what we just sent. This is the plain-text-with-tool-parser path.
            buffered_fallback = (
                has_tool_parser
                and not (native_streaming and not native_streaming_error)
            )
            if buffered_fallback and not tool_calls_proto and content:
                yield backend_pb2.Reply(
                    message=bytes(content, encoding='utf-8'),
                    chat_deltas=[backend_pb2.ChatDelta(content=content)],
                )
                chat_delta = backend_pb2.ChatDelta(
                    reasoning_content=reasoning_content,
                    tool_calls=tool_calls_proto,
                )
            yield backend_pb2.Reply(
                message=b"",
                prompt_tokens=prompt_tokens,
                tokens=completion_tokens,
                chat_deltas=[chat_delta],
                logprobs=logprobs_bytes,
            )
            return

        # Non-streaming: single Reply with everything
        yield backend_pb2.Reply(
            message=bytes(content, encoding='utf-8'),
            prompt_tokens=prompt_tokens,
            tokens=completion_tokens,
            chat_deltas=[chat_delta],
            logprobs=logprobs_bytes,
        )

    def load_image(self, image_path: str):
        """
        Load an image from the given file path or base64 encoded data.

        Args:
            image_path (str): The path to the image file or base64 encoded data.

        Returns:
            Image: The loaded image.
        """
        try:

            image_data = base64.b64decode(image_path)
            image = Image.open(io.BytesIO(image_data))
            return image
        except Exception as e:
            print(f"Error loading image {image_path}: {e}", file=sys.stderr)
            return None

    def load_video(self, video_path: str):
        """
        Load a video from the given file path.

        Args:
            video_path (str): The path to the image file.

        Returns:
            Video: The loaded video.
        """
        try:
            timestamp = str(int(time.time() * 1000))  # Generate timestamp
            p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data")
            with open(p, "wb") as f:
                f.write(base64.b64decode(video_path))
            video = VideoAsset(name=p).np_ndarrays
            os.remove(p)
            return video
        except Exception as e:
            print(f"Error loading video {video_path}: {e}", file=sys.stderr)
            return None

async def serve(address):
    # Start asyncio gRPC server
    server = grpc.aio.server(migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
        options=[
            ('grpc.max_message_length', 50 * 1024 * 1024),  # 50MB
            ('grpc.max_send_message_length', 50 * 1024 * 1024),  # 50MB
            ('grpc.max_receive_message_length', 50 * 1024 * 1024),  # 50MB
        ],
        interceptors=get_auth_interceptors(aio=True),
    )
    # Add the servicer to the server
    backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
    # Bind the server to the address
    server.add_insecure_port(address)

    # Gracefully shutdown the server on SIGTERM or SIGINT
    loop = asyncio.get_event_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(
            sig, lambda: asyncio.ensure_future(server.stop(5))
        )

    # Start the server
    await server.start()
    print("Server started. Listening on: " + address, file=sys.stderr)
    # Wait for the server to be terminated
    await server.wait_for_termination()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the gRPC server.")
    parser.add_argument(
        "--addr", default="localhost:50051", help="The address to bind the server to."
    )
    args = parser.parse_args()

    asyncio.run(serve(args.addr))
