# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# 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.
"""
Handler for the /v1/chat/completions endpoint.

Supports streaming (SSE via DirectStreamer) and non-streaming (JSON) responses.
"""

import asyncio
import time
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING

from ...utils import logging
from ...utils.import_utils import is_serve_available


if is_serve_available():
    from fastapi.responses import JSONResponse, StreamingResponse
    from openai.types.chat import (
        ChatCompletion,
        ChatCompletionMessage,
        ChatCompletionMessageToolCall,
    )
    from openai.types.chat.chat_completion import Choice
    from openai.types.chat.chat_completion_chunk import (
        ChatCompletionChunk,
        ChoiceDelta,
        ChoiceDeltaToolCall,
    )
    from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
    from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
    from openai.types.completion_usage import CompletionUsage


from .utils import (
    BaseGenerateManager,
    BaseHandler,
    Modality,
    ReasoningText,
    _StreamError,
    get_reasoning_config,
    get_tool_call_config,
    parse_reasoning,
    parse_tool_calls,
)


if TYPE_CHECKING:
    from transformers import GenerationConfig, PreTrainedModel, PreTrainedTokenizerFast, ProcessorMixin


class TransformersCompletionCreateParamsStreaming(CompletionCreateParamsStreaming, total=False):
    generation_config: str
    seed: int
    chat_template_kwargs: dict


# Fields accepted by the OpenAI schema but not yet supported.
# Receiving these raises an error to avoid silent misbehaviour.
# NOTE: "stop" is NOT in this set — we map it to stop_strings.
UNUSED_CHAT_COMPLETION_FIELDS = {
    "audio",
    "function_call",
    "functions",
    "logprobs",
    "max_completion_tokens",
    "metadata",
    "modalities",
    "n",
    "parallel_tool_calls",
    "prediction",
    "presence_penalty",
    "reasoning_effort",
    "response_format",
    "service_tier",
    "store",
    "stream_options",
    "tool_choice",
    "top_logprobs",
    "user",
    "web_search_options",
}


logger = logging.get_logger(__name__)


class ChatCompletionHandler(BaseHandler):
    """Handler for the `/v1/chat/completions` endpoint.

    Supports both streaming (SSE) and non-streaming (JSON) responses.
    """

    _valid_params_class = TransformersCompletionCreateParamsStreaming
    _unused_fields = UNUSED_CHAT_COMPLETION_FIELDS

    async def handle_request(self, body: dict, request_id: str) -> StreamingResponse | JSONResponse:
        """Validate the request, load the model, and dispatch to streaming or non-streaming.

        Args:
            body (`dict`): The raw JSON request body (OpenAI chat completion format).
            request_id (`str`): Unique request identifier (from header or auto-generated).

        Returns:
            `StreamingResponse | JSONResponse`: SSE stream or JSON depending on ``body["stream"]``.
        """
        self._validate_request(body)

        model_id, model, processor = self._resolve_model(body)
        modality = self.model_manager.get_model_modality(model, processor=processor)
        use_cb = self.generation_state.use_continuous_batching(model, modality)
        logger.warning(f"[Request received] Model: {model_id}, CB: {use_cb}")
        gen_manager = self.generation_state.get_manager(model_id, use_cb=use_cb)
        processor_inputs = self.get_processor_inputs_from_messages(body["messages"], modality)

        has_video = any(
            c.get("type") == "video"
            for msg in processor_inputs
            for c in (msg.get("content") if isinstance(msg.get("content"), list) else [])
        )
        # Default to 32 frames for video (Gemma 4 default); some processors load all frames otherwise.
        # Merge order (later wins): custom default -> server default → request-level kwargs.
        chat_template_kwargs: dict = {}
        if has_video:
            chat_template_kwargs["num_frames"] = 32
        chat_template_kwargs.update(self.chat_template_kwargs)
        chat_template_kwargs.update(body.get("chat_template_kwargs", {}))
        inputs = processor.apply_chat_template(
            processor_inputs,
            add_generation_prompt=True,
            tools=body.get("tools"),
            return_tensors=None if use_cb else "pt",
            return_dict=True,
            tokenize=True,
            load_audio_from_video=modality == Modality.MULTIMODAL and has_video,
            **chat_template_kwargs,
        )
        if not use_cb:
            inputs = inputs.to(model.device)  # type: ignore[union-attr]

        gen_config = self._build_generation_config(body, model.generation_config, use_cb=use_cb)
        # TODO: remove when CB supports per-request generation config
        if use_cb:
            gen_manager.init_cb(model, gen_config)

        tool_config = get_tool_call_config(processor, model) if body.get("tools") else None
        reasoning_config = get_reasoning_config(processor, model, inputs["input_ids"])

        streaming = body.get("stream")
        if streaming:
            return self._streaming(
                request_id,
                model,
                processor,
                model_id,
                inputs,
                gen_config,
                gen_manager=gen_manager,
                tool_config=tool_config,
                reasoning_config=reasoning_config,
            )
        else:
            return await self._non_streaming(
                request_id,
                model,
                processor,
                model_id,
                inputs,
                gen_config,
                gen_manager=gen_manager,
                tool_config=tool_config,
                reasoning_config=reasoning_config,
            )

    # ----- streaming -----

    def _streaming(
        self,
        request_id: str,
        model: "PreTrainedModel",
        processor: "ProcessorMixin | PreTrainedTokenizerFast",
        model_id: str,
        inputs: dict,
        gen_config: "GenerationConfig",
        gen_manager: BaseGenerateManager,
        tool_config: dict | None = None,
        reasoning_config: dict | None = None,
    ) -> StreamingResponse:
        """Stream tokens as SSE via DirectStreamer."""
        queue, streamer = gen_manager.generate_streaming(
            model,
            processor,
            inputs,
            gen_config,
            request_id=request_id,
            tool_config=tool_config,
            reasoning_config=reasoning_config,
        )
        input_ids = inputs["input_ids"]
        # CB returns plain lists, regular path returns tensors
        input_len = len(input_ids) if isinstance(input_ids, list) else input_ids.shape[-1]

        async def sse_gen() -> AsyncGenerator[str, None]:
            try:
                yield self._build_chunk_sse(request_id, role="assistant", model=model_id)

                done = False
                while not done:
                    text = await queue.get()
                    batch = [text]
                    try:
                        while True:
                            batch.append(queue.get_nowait())
                    except asyncio.QueueEmpty:
                        pass

                    sse_parts: list[str] = []
                    for text in batch:
                        if text is None:
                            done = True
                            break
                        if isinstance(text, _StreamError):
                            sse_parts.append(f'data: {{"error": "{text.msg}"}}\n\n')
                            yield "".join(sse_parts)
                            return

                        if isinstance(text, ReasoningText):
                            sse_parts.append(self._build_chunk_sse(request_id, model=model_id, reasoning_content=text))
                        else:
                            sse_parts.append(self._build_chunk_sse(request_id, model=model_id, content=text))

                    if sse_parts:
                        yield "".join(sse_parts)

                # Tool calls are parsed after generation completes (not during streaming),
                # because the full token sequence is needed for reliable parsing.
                has_tool_calls = False
                if tool_config:
                    parsed = parse_tool_calls(processor, streamer.generated_token_ids, tool_config["schema"])
                    if parsed:
                        has_tool_calls = True
                        for i, tc in enumerate(parsed):
                            yield self._build_chunk_sse(
                                request_id,
                                model=model_id,
                                tool_calls=[
                                    ChoiceDeltaToolCall(
                                        index=i,
                                        type="function",
                                        id=f"{request_id}_tool_call_{i}",
                                        function={"name": tc["name"], "arguments": tc["arguments"]},
                                    )
                                ],
                            )

                hit_max = gen_config.max_new_tokens is not None and streamer.total_tokens >= gen_config.max_new_tokens
                if has_tool_calls:
                    finish_reason = "tool_calls"
                elif hit_max:
                    finish_reason = "length"
                else:
                    finish_reason = "stop"
                usage = CompletionUsage(
                    prompt_tokens=input_len,
                    completion_tokens=streamer.total_tokens,
                    total_tokens=input_len + streamer.total_tokens,
                )
                yield self._build_chunk_sse(
                    request_id,
                    finish_reason=finish_reason,
                    model=model_id,
                    usage=usage,
                )
            except (GeneratorExit, asyncio.CancelledError):
                # Client disconnected — abort generation to free GPU.
                # Re-raise is mandatory: Python raises RuntimeError if GeneratorExit is swallowed.
                streamer.cancel()
                raise

        return StreamingResponse(sse_gen(), media_type="text/event-stream")

    # ----- non-streaming -----

    async def _non_streaming(
        self,
        request_id: str,
        model: "PreTrainedModel",
        processor: "ProcessorMixin | PreTrainedTokenizerFast",
        model_id: str,
        inputs: dict,
        gen_config: "GenerationConfig",
        gen_manager: BaseGenerateManager,
        tool_config: dict | None = None,
        reasoning_config: dict | None = None,
    ) -> JSONResponse:
        """Run generation and return a JSONResponse."""
        content, input_len, generated_ids = await gen_manager.generate_non_streaming(
            model, processor, inputs, gen_config, request_id=request_id
        )

        hit_max = gen_config.max_new_tokens is not None and len(generated_ids) >= gen_config.max_new_tokens
        completion_tokens = len(generated_ids)
        usage = CompletionUsage(
            prompt_tokens=input_len,
            completion_tokens=completion_tokens,
            total_tokens=input_len + completion_tokens,
        )

        tool_calls = None
        if tool_config is not None:
            parsed = parse_tool_calls(processor, generated_ids, tool_config["schema"])
            if parsed:
                tool_calls = [
                    ChatCompletionMessageToolCall(
                        id=f"{request_id}_tool_call_{i}",
                        type="function",
                        function={"name": tc["name"], "arguments": tc["arguments"]},
                    )
                    for i, tc in enumerate(parsed)
                ]

        reasoning_content = None
        if reasoning_config is not None:
            content, reasoning_content = parse_reasoning(processor, generated_ids, content, reasoning_config)

        if tool_calls is not None:
            finish_reason = "tool_calls"
        elif hit_max:
            finish_reason = "length"
        else:
            finish_reason = "stop"

        return JSONResponse(
            self._build_completion(
                request_id,
                content,
                model_id,
                finish_reason=finish_reason,
                usage=usage,
                tool_calls=tool_calls,
                reasoning_content=reasoning_content,
            ),
            media_type="application/json",
        )

    # ----- helpers -----

    def _build_generation_config(self, body: dict, model_generation_config: "GenerationConfig", use_cb: bool = False):
        """Apply Chat Completions params (``max_tokens``, ``frequency_penalty``, ``logit_bias``,
        ``stop``) on top of the base generation config."""
        generation_config = super()._build_generation_config(body, model_generation_config, use_cb=use_cb)

        if body.get("max_tokens") is not None:
            generation_config.max_new_tokens = int(body["max_tokens"])
        if body.get("frequency_penalty") is not None:
            generation_config.repetition_penalty = 1.0 + float(body["frequency_penalty"])
        if body.get("logit_bias") is not None:
            generation_config.sequence_bias = {(int(k),): v for k, v in body["logit_bias"].items()}
        if body.get("stop") is not None:
            generation_config.stop_strings = body["stop"]

        return generation_config

    # ----- response builders -----

    def _build_completion(
        self,
        request_id: str,
        content: str,
        model_id: str,
        finish_reason: str,
        usage: CompletionUsage | None = None,
        tool_calls: list[dict] | None = None,
        reasoning_content: str | None = None,
    ) -> dict:
        """Build a non-streaming ChatCompletion response dict.

        Args:
            request_id (`str`): Unique request identifier.
            content (`str`): The generated text.
            model_id (`str`): Model ID to include in the response.
            finish_reason (`str`): Why generation stopped (``"stop"``, ``"length"``, ``"tool_calls"``).
            usage (`CompletionUsage`, *optional*): Token usage statistics.
            tool_calls (`list[dict]`, *optional*): Parsed tool calls, if any.
            reasoning_content (`str`, *optional*): Chain-of-thought content extracted from the response.

        Returns:
            `dict`: Serialized ``ChatCompletion`` ready for JSON response.
        """
        # reasoning_content is added as an extra field (base types set extra="allow")
        # we use model_validate rather than __init__ to avoid ty raising errors for the extra field
        message = ChatCompletionMessage.model_validate(
            {"content": content, "role": "assistant", "tool_calls": tool_calls, "reasoning_content": reasoning_content}
        )
        result = ChatCompletion(
            id=request_id,
            created=int(time.time()),
            object="chat.completion",
            model=model_id,
            choices=[Choice(index=0, message=message, finish_reason=finish_reason)],
            usage=usage,
        )
        return result.model_dump(exclude_none=True)

    def _build_chunk_sse(
        self,
        request_id: str = "",
        content: str | None = None,
        model: str | None = None,
        role: str | None = None,
        finish_reason: str | None = None,
        tool_calls: list | None = None,
        usage: CompletionUsage | None = None,
        reasoning_content: str | None = None,
    ) -> str:
        """Build a streaming ``ChatCompletionChunk`` and format it as an SSE ``data:`` line.

        Args:
            request_id (`str`): Unique request identifier.
            content (`str`, *optional*): Text content delta.
            model (`str`, *optional*): Model ID.
            role (`str`, *optional*): Role (only sent in the first chunk).
            finish_reason (`str`, *optional*): Set on the final chunk.
            tool_calls (`list`, *optional*): Tool call deltas.
            usage (`CompletionUsage`, *optional*): Token usage (sent with the final chunk).
            reasoning_content (`str`, *optional*): Reasoning/thinking delta (OpenAI-compatible extension).

        Returns:
            `str`: A formatted SSE event string.
        """
        # reasoning_content is added as an extra field (base types set extra="allow")
        # we use model_validate rather than __init__ to avoid ty raising errors for the extra field
        delta = ChoiceDelta.model_validate(
            {"content": content, "role": role, "tool_calls": tool_calls, "reasoning_content": reasoning_content}
        )
        chunk = ChatCompletionChunk(
            id=request_id,
            created=int(time.time()),
            model=model,
            choices=[ChoiceChunk(delta=delta, index=0, finish_reason=finish_reason)],
            usage=usage,
            system_fingerprint="",
            object="chat.completion.chunk",
        )
        return self.chunk_to_sse(chunk)
