import sys
import asyncio
import statistics
import time
from threading import Thread, Event
import logging
import platform
import tempfile
import traceback
from typing import Optional, Union, List
import json
from pathlib import Path
import os
import shutil
from fastapi import FastAPI, HTTPException, status, Request, WebSocket, Form, UploadFile
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from starlette.websockets import WebSocketDisconnect, WebSocketState
import uvicorn
from uvicorn.config import Config
from uvicorn.server import Server as UvicornServer
from tabulate import tabulate

from openai.types.completion import Completion, CompletionChoice
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from openai.types.chat import ChatCompletionMessage
from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
    Function,
)
from openai.types.completion_usage import CompletionUsage
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import (
    ChoiceDelta,
    ChoiceDeltaToolCall,
    ChoiceDeltaToolCallFunction,
)
from openai.types.completion_choice import Logprobs
from openai.types.model import Model
from openai.types.responses import (
    Response,
    ResponseOutputMessage,
    ResponseOutputText,
    ResponseCreatedEvent,
    ResponseTextDeltaEvent,
    ResponseCompletedEvent,
)

import lemonade.api as lemonade_api
from lemonade.tools.server.wrapped_server import WrappedServer
from lemonade.tools.server.llamacpp import LlamaServer
from lemonade.tools.server.flm import FlmServer
from lemonade.tools.server.tool_calls import extract_tool_calls, get_tool_call_pattern
from lemonade.tools.server.webapp import get_webapp_html
from lemonade.tools.server.utils.port import lifespan

from lemonade_server.model_manager import ModelManager
from lemonade_server.pydantic_models import (
    DEFAULT_PORT,
    DEFAULT_HOST,
    DEFAULT_LOG_LEVEL,
    DEFAULT_LLAMACPP_BACKEND,
    DEFAULT_CTX_SIZE,
    LoadConfig,
    CompletionRequest,
    ChatCompletionRequest,
    EmbeddingsRequest,
    RerankingRequest,
    ResponsesRequest,
    PullConfig,
    DeleteConfig,
    LogLevelConfig,
)
from lemonade_server.settings import save_setting

# Set to a high number to allow for interesting experiences in real apps
# Tests should use the max_new_tokens argument to set a lower value
DEFAULT_MAX_NEW_TOKENS = 1500

if platform.system() in ["Windows", "Darwin"]:
    # pylint: disable=ungrouped-imports
    from lemonade.tools.server.tray import LemonadeTray, OutputDuplicator


class ServerLogFilter(logging.Filter):
    def __init__(self, server):
        super().__init__()
        self.server = server
        self.noisy_paths = {
            "/api/v1/health",
            "/api/v0/health",
            "/api/v1/models",
            "/api/v0/models",
        }

    def filter(self, record: logging.LogRecord) -> bool:
        msg = record.getMessage()

        # Filter out websocket logs
        if "> TEXT" in msg:
            return False

        # Filter out noisy HTTP routes if debug logs are OFF
        if not self.server.debug_logging_enabled:
            if any(path in msg for path in self.noisy_paths):
                return False

        # Otherwise, allow the log
        return True


async def log_streamer(websocket: WebSocket, path: str, interval: float = 1.0):
    logger = logging.getLogger()
    await websocket.accept()
    try:
        with open(path, "r", encoding="utf-8") as f:
            f.seek(0)  # start at the beginning of the file
            while True:
                # Try reading a line
                line = f.readline()
                if not line:
                    await asyncio.sleep(interval)
                    continue

                # Send defensively: if disconnected, bail out
                if websocket.application_state != WebSocketState.CONNECTED:
                    # Server-side state says we're not connected anymore
                    break

                try:
                    await websocket.send_text(line)
                except WebSocketDisconnect:
                    # Client closed — normal path out
                    break
                except RuntimeError as re:
                    # Starlette will raise this if a close has already been sent
                    logger.debug("RuntimeError during send: %s", re)
                    break

    except WebSocketDisconnect:
        # Client closed the socket; do not try to send or close again
        pass
    except Exception as e:  # pylint: disable=broad-except
        # Log server-side; do not attempt to send error over a possibly closed socket
        logger.exception("Error in log_streamer: %s", e)
    finally:
        # Only close if Starlette still thinks we're connected.
        # This prevents "Cannot call send once a close message has been sent."
        try:
            if websocket.application_state == WebSocketState.CONNECTED:
                await websocket.close()
        except Exception:  # pylint: disable=broad-except
            # If close itself races, swallow — we're shutting down anyway.
            pass


class ServerModel(Model):
    """
    An extension of OpenAI's Model class that adds
    checkpoint and recipe attributes.
    """

    checkpoint: str
    recipe: str


class GeneratorThread(Thread):
    """
    Thread class designed for use with streaming generation within
    an LLM server. It needs access to the streamer in order to order
    to help the completions APIs escape the "for text in streamer" loop.
    It also provides exception handling that works nicely with HTTP
    servers by providing the stack trace and making the exception
    information available to the main thread.
    """

    def __init__(self, streamer, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.exception = None
        self.streamer = streamer

    def run(self):
        try:
            if self._target:
                self._target(*self._args, **self._kwargs)
        except Exception as e:  # pylint: disable=broad-except
            self.exception = e
            logging.error(f"Exception raised in generate thread: {e}")
            traceback.print_exc()
            self.streamer.done()


class StopOnEvent:
    """
    Custom stopping criteria that halts text generation when a specified event is set.

    This allows for external control of generation, such as stopping a generation
    before it reaches the maximum token limit.
    """

    def __init__(self, stop_event: Event):
        super().__init__()
        self.stop_event = stop_event

    def __call__(self, input_ids, scores, **kwargs):
        return self.stop_event.is_set()


class NoCacheStaticFiles(StaticFiles):
    """Custom StaticFiles class with no-cache headers"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def file_response(self, *args, **kwargs) -> Response:
        response = super().file_response(*args, **kwargs)
        # Add no-cache headers for all static files
        response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
        response.headers["Pragma"] = "no-cache"
        response.headers["Expires"] = "0"
        return response


class Server:
    """
    Open a web server that apps can use to communicate with the LLM.

    The server exposes these endpoints:
    - /api/v1/pull: install an LLM by its Lemonade Server Model Name.
    - /api/v1/delete: delete an LLM by its Lemonade Server Model Name.
    - /api/v1/load: load a model checkpoint.
    - /api/v1/unload: unload a model checkpoint.
    - /api/v1/health: check whether a model is loaded and ready to serve.
    - /api/v1/stats: performance statistics for the generation.
    - /api/v1/halt: stop an in-progress generation from make more tokens.
    - /api/v1/completions: completion responses using HTTP chunked transfer encoding.
    - /api/v1/chat/completions: chat completion responses using HTTP chunked transfer encoding.
    - /api/v1/responses: responses API using HTTP chunked transfer encoding.
    - /api/v1/models: list all available models.
    - /api/v1/models/{model_id}: retrieve a specific model by ID.
    """

    def __init__(
        self,
        port: int = DEFAULT_PORT,
        host: str = DEFAULT_HOST,
        log_level: str = DEFAULT_LOG_LEVEL,
        ctx_size: int = DEFAULT_CTX_SIZE,
        tray: bool = False,
        log_file: str = None,
        llamacpp_backend: str = DEFAULT_LLAMACPP_BACKEND,
    ):
        super().__init__()

        # Save args as members
        self.port = port
        self.host = host
        self.log_level = log_level
        self.ctx_size = ctx_size
        self.tray = tray
        self.log_file = log_file
        self.llamacpp_backend = llamacpp_backend

        # Initialize FastAPI app
        self.app = FastAPI(lifespan=lifespan)

        # Lifespan will load some tasks in the background, and then set the
        # app.initialized flag to True when this is done
        self.app.initialized = False

        # Add CORS middleware
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],  # Allows all origins
            allow_credentials=True,
            allow_methods=["*"],  # Allows all methods
            allow_headers=["*"],  # Allows all headers
        )

        # Set up debug middleware if debug logging is enabled
        # This must be done during app initialization, not at runtime
        self.debug_logging_enabled = log_level == "debug"
        if self.debug_logging_enabled:
            self.setup_middleware_timer()

        # Set up custom routes
        self.setup_routes(["/api/v0", "/api/v1"])

        # Set up Web App
        self.app.get("/")(self.webapp)

        # Mount a static assets dir for HTML responses, such
        # as the Web App
        static_dir = Path(__file__).parent / "static"
        self.app.mount(
            "/static", NoCacheStaticFiles(directory=static_dir), name="static_assets"
        )

        # Performance stats that are set during /ws and can be
        # fetched in /stats
        self.time_to_first_token = None
        self.tokens_per_second = None
        self.input_tokens = None
        self.output_tokens = None
        self.decode_token_times = None

        # Store debug logging state
        self.debug_logging_enabled = logging.getLogger().isEnabledFor(logging.DEBUG)

        # Flag that tells the LLM to stop generating text and end the response
        self.stop_event = Event()

        self.llm_loaded: LoadConfig = None
        self.tokenizer = None

        # Placeholders for model and configs
        self.model = None

        # Initialize semaphore for tracking active generations
        self.max_concurrent_generations = 1
        self._generate_semaphore = asyncio.Semaphore(self.max_concurrent_generations)

        # Dictionary of installed LLM, by model name : information about those models
        # Does not include non-installed models
        self.local_models = ModelManager().downloaded_models_enabled

        # Add lock for load/unload operations
        self._load_lock = asyncio.Lock()

        # Subprocess handle for wrapped instance of llama_server.exe, etc.
        self.wrapped_server: WrappedServer = None

    def setup_routes(self, api_prefixes: list[str]):
        for prefix in api_prefixes:
            # Custom routes
            self.app.post(f"{prefix}/pull")(self.pull)
            self.app.post(f"{prefix}/delete")(self.delete)
            self.app.post(f"{prefix}/load")(self.load_llm)
            self.app.post(f"{prefix}/unload")(self.unload_llm)
            self.app.get(f"{prefix}/health")(self.health)
            self.app.get(f"{prefix}/halt")(self.halt_generation)
            self.app.get(f"{prefix}/stats")(self.send_stats)
            self.app.get(f"{prefix}/system-info")(self.get_system_info)
            self.app.post(f"{prefix}/completions")(self.completions)
            self.app.post(f"{prefix}/responses")(self.responses)
            self.app.post(f"{prefix}/log-level")(self.set_log_level)
            self.app.websocket(f"{prefix}/logs/ws")(self.logs_ws)
            self.app.post(f"{prefix}/add-local-model")(self.add_local_model)

            # OpenAI-compatible routes
            self.app.post(f"{prefix}/chat/completions")(self.chat_completions)
            self.app.post(f"{prefix}/embeddings")(self.embeddings)
            self.app.get(f"{prefix}/models")(self.models)
            self.app.get(f"{prefix}/models/{{model_id}}")(self.retrieve_model)

            # JinaAI routes (jina.ai/reranker/)
            self.app.post(f"{prefix}/reranking")(self.reranking)
            self.app.post(f"{prefix}/rerank")(self.reranking)

            # Migration routes
            self.app.get(f"{prefix}/migration/incompatible-models")(
                self.get_incompatible_models
            )
            self.app.post(f"{prefix}/migration/cleanup")(
                self.cleanup_incompatible_models
            )

    async def add_local_model(
        self,
        model_name: str = Form(...),
        checkpoint: str = Form(""),
        recipe: str = Form(...),
        reasoning: bool = Form(False),
        vision: bool = Form(False),
        mmproj: str = Form(None),
        model_files: List[UploadFile] = None,
    ):
        from huggingface_hub.constants import HF_HUB_CACHE
        from lemonade.tools.llamacpp.utils import parse_checkpoint

        # Upload and register a local model from files.
        try:
            if not model_files:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="No model files provided for upload",
                )

            if not model_name.startswith("user."):
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="Model name must start with 'user.'",
                )

            valid_recipes = ["llamacpp", "oga-npu", "oga-hybrid", "oga-cpu"]
            if recipe not in valid_recipes:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=f"Invalid recipe. Must be one of: {', '.join(valid_recipes)}",
                )

            if recipe == "llamacpp" and not any(
                f.filename.lower().endswith(".gguf") for f in model_files
            ):
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="At least one .gguf file is required for llamacpp",
                )

            # Check if model name already exists
            if model_name in ModelManager().supported_models:
                raise HTTPException(
                    status_code=status.HTTP_409_CONFLICT,
                    detail=(
                        f"Model name '{model_name}' already exists. "
                        "Please use a different name."
                    ),
                )

            model_name_clean = model_name.replace("user.", "")

            # Files are saved to models--{model_name_clean}
            # Note: This is based on the user's custom model name, NOT the checkpoint field
            repo_cache_name = model_name_clean.replace("/", "--")
            snapshot_path = os.path.join(HF_HUB_CACHE, f"models--{repo_cache_name}")
            os.makedirs(snapshot_path, exist_ok=True)

            # Extract variant from checkpoint field if provided
            # checkpoint field format: "folder:variant" or just "folder"
            variant = None
            if checkpoint and ":" in checkpoint:
                _, variant = parse_checkpoint(checkpoint)
                # variant now contains just the variant[can be with or without the
                # .gguf extension] filename (e.g., "LFM2-VL-1.6B-F16 or LFM2-VL-1.6B-F16.gguf")

            # Save uploaded files, preserving folder structure
            for file in model_files:
                relative_path = file.filename
                path_parts = relative_path.split("/")

                if len(path_parts) > 1:
                    internal_path = "/".join(path_parts[1:])
                    file_path = os.path.join(snapshot_path, internal_path)
                else:
                    file_path = os.path.join(snapshot_path, path_parts[0])

                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                with open(file_path, "wb") as f:
                    content = await file.read()
                    f.write(content)

            # Resolve actual file paths after upload (for faster loading later)
            resolved_checkpoint = None
            resolved_mmproj = None

            # For OGA models, find genai_config.json
            if recipe.startswith("oga-"):
                for root, _, files in os.walk(snapshot_path):
                    if "genai_config.json" in files:
                        resolved_checkpoint = root
                        break
                if not resolved_checkpoint:
                    resolved_checkpoint = snapshot_path

            # For llamacpp models, find the GGUF file
            elif recipe == "llamacpp":
                gguf_file_found = None

                # If variant is specified, look for that specific file
                if variant:
                    search_term = (
                        variant if variant.endswith(".gguf") else f"{variant}.gguf"
                    )
                    for root, _, files in os.walk(snapshot_path):
                        if search_term in files:
                            gguf_file_found = os.path.join(root, search_term)
                            break

                # If no variant or variant not found, search for any .gguf file (excluding mmproj)
                if not gguf_file_found:
                    for root, _, files in os.walk(snapshot_path):
                        gguf_files = [
                            f
                            for f in files
                            if f.endswith(".gguf") and "mmproj" not in f.lower()
                        ]
                        if gguf_files:
                            gguf_file_found = os.path.join(root, gguf_files[0])
                            break

                resolved_checkpoint = (
                    gguf_file_found if gguf_file_found else snapshot_path
                )

            # Search for mmproj file if provided
            if mmproj:
                for root, _, files in os.walk(snapshot_path):
                    if mmproj in files:
                        resolved_mmproj = os.path.join(root, mmproj)
                        break

            # Build checkpoint for registration
            # For llamacpp with resolved path, store the full path relative to HF_HUB_CACHE
            if resolved_checkpoint:
                # Store as relative path from HF_HUB_CACHE for portability
                checkpoint_to_register = os.path.relpath(
                    resolved_checkpoint, HF_HUB_CACHE
                )
            elif variant:
                checkpoint_to_register = f"models--{repo_cache_name}:{variant}"
            else:
                checkpoint_to_register = f"models--{repo_cache_name}"

            # Register the model
            ModelManager().register_local_model(
                model_name=model_name,
                checkpoint=checkpoint_to_register,
                recipe=recipe,
                reasoning=reasoning,
                vision=vision,
                mmproj=resolved_mmproj if resolved_mmproj else mmproj,
                snapshot_path=snapshot_path,
            )

            # Refresh local models
            self.local_models = ModelManager().downloaded_models_enabled

            return {
                "status": "success",
                "message": f"Model {model_name} uploaded and registered successfully",
            }
        except Exception as e:
            if os.path.exists(checkpoint_to_register):
                shutil.rmtree(checkpoint_to_register)
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Failed to upload model: {str(e)}",
            )

    async def set_log_level(self, config: LogLevelConfig):
        """
        Set the logging level of the server.
        """
        try:
            log_level_upper = config.level.upper()
            numeric_level = getattr(logging, log_level_upper, None)
            if not isinstance(numeric_level, int):
                raise ValueError(f"Invalid log level: {config.level}")

            # Get the root logger
            logger = logging.getLogger()
            logger.setLevel(numeric_level)

            # Update all handlers
            for handler in logger.handlers:
                handler.setLevel(numeric_level)

            logging.getLogger("uvicorn.error").setLevel(numeric_level)
            self.debug_logging_enabled = numeric_level <= logging.DEBUG

            # Save the setting
            save_setting("log_level", config.level)

            logging.info(f"Log level changed to: {config.level}")
            return {"status": "success", "message": f"Log level set to {config.level}"}
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Failed to set log level: {str(e)}",
            )

    def _log_request_parameters(self, request, endpoint_name: str):
        """
        Log request parameters excluding content fields like messages, prompt, or input.

        Args:
            request: Any request object (CompletionRequest, ChatCompletionRequest, etc.)
            endpoint_name: Name of the endpoint for logging context
        """
        if not logging.getLogger().isEnabledFor(logging.DEBUG):
            return

        # Fields to exclude from logging (content fields)
        excluded_fields = {"messages", "prompt", "input"}

        # Get all attributes from the request object
        request_params = {}
        if hasattr(request, "__dict__"):
            # For pydantic models, get the dict representation
            if hasattr(request, "model_dump"):
                all_params = request.model_dump()
            elif hasattr(request, "dict"):
                all_params = request.dict()
            else:
                all_params = request.__dict__

            # Filter out excluded fields and add special handling for certain fields
            for key, value in all_params.items():
                if key not in excluded_fields:
                    # Special handling for tools field - show count instead of full content
                    if key == "tools" and value is not None:
                        request_params[key] = (
                            f"{len(value)} tools" if isinstance(value, list) else value
                        )
                    # Special handling for input type in responses
                    elif key == "input" and hasattr(request, "input"):
                        request_params["input_type"] = type(value).__name__
                    else:
                        request_params[key] = value

        logging.debug(f"{endpoint_name} request parameters: {request_params}")

    def _setup_server_common(
        self,
        tray: bool = False,
        threaded_mode: bool = False,
    ):
        """
        Common setup logic shared between run() and run_in_thread().

        Args:
            tray: Whether to run the server in tray mode
            threaded_mode: Whether this is being set up for threaded execution
        """

        # Define TRACE level
        logging.TRACE = 9  # Lower than DEBUG which is 10
        logging.addLevelName(logging.TRACE, "TRACE")

        # Add a convenience function at the module level
        def trace(message, *args, **kwargs):
            logging.log(logging.TRACE, message, *args, **kwargs)

        logging.trace = trace

        # Configure logging based on mode
        if threaded_mode:
            # Configure logging for warning level (to reduce noise in threaded execution)
            logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
        else:
            # Configure logging to match uvicorn's format
            logging_level = getattr(logging, self.log_level.upper())

            # Set up file handler for logging to lemonade.log
            uvicorn_formatter = uvicorn.logging.DefaultFormatter(
                fmt="%(levelprefix)s %(message)s",
                use_colors=True,
            )
            if not self.log_file:
                self.log_file = tempfile.NamedTemporaryFile(
                    prefix="lemonade_", suffix=".log", delete=False
                ).name
            file_handler = logging.FileHandler(
                self.log_file, mode="a", encoding="utf-8"
            )
            file_handler.setLevel(logging_level)
            file_handler.setFormatter(uvicorn_formatter)
            file_handler.addFilter(ServerLogFilter(self))

            # Set up console handler
            console_handler = logging.StreamHandler()
            console_handler.setLevel(logging_level)
            console_handler.setFormatter(uvicorn_formatter)
            console_handler.addFilter(ServerLogFilter(self))

            # Configure root logger with both handlers
            logging.basicConfig(
                level=logging_level,
                handlers=[file_handler, console_handler],
                force=True,
            )

        # Update debug logging state after setting log level
        self.debug_logging_enabled = logging.getLogger().isEnabledFor(logging.DEBUG)
        if tray:
            # Save original stdout/stderr
            sys.stdout = OutputDuplicator(self.log_file, sys.stdout)
            sys.stderr = OutputDuplicator(self.log_file, sys.stderr)

            # Open lemonade server in tray mode
            # lambda function used for deferred instantiation and thread safety
            LemonadeTray(
                self.log_file, self.port, lambda: self, log_level=self.log_level
            ).run()
            sys.exit(0)

        # Let the app know what port it's running on, so
        # that the lifespan can access it
        self.app.port = self.port
        # FastAPI already has a `host` function and we cannot use `_host` as
        # PyLint will believe its private
        self.app.host_ = self.host

    def run(self):
        # Common setup
        self._setup_server_common(
            threaded_mode=False,
            tray=self.tray,
        )

        uvicorn.run(self.app, host=self.host, port=self.port, log_level=self.log_level)

    def run_in_thread(self, host: str = "localhost"):
        """
        Set up the server for running in a thread.
        Returns a uvicorn server instance that can be controlled externally.
        """
        # Common setup
        self._setup_server_common(
            threaded_mode=True,
            tray=False,
        )

        class CustomServer(UvicornServer):
            """Custom Uvicorn server that can be properly shutdown from another thread"""

            def install_signal_handlers(self):
                pass

        # Configure the server
        config = Config(
            app=self.app,
            host=host,
            port=self.port,
            log_level=self.log_level,
            log_config=None,
        )

        # Create and return the uvicorn server
        return CustomServer(config=config)

    async def _show_telemetry(self):
        """
        Show telemetry data in debug mode.
        """
        # Exit early if debug logging is disabled or no telemetry data is available
        if not self.debug_logging_enabled or self.tokens_per_second is None:
            return

        # Prepare telemetry data (transposed format)
        telemetry = [
            ["Input tokens", self.input_tokens],
            ["Output tokens", self.output_tokens],
            ["TTFT (s)", f"{self.time_to_first_token:.2f}"],
            ["TPS", f"{self.tokens_per_second:.2f}"],
        ]

        table = tabulate(
            telemetry, headers=["Metric", "Value"], tablefmt="fancy_grid"
        ).split("\n")

        # Show telemetry in debug while complying with uvicorn's log indentation
        logging.debug("\n          ".join(table))

    def webapp(self):
        """
        Serve the Web App to the user's browser.
        """

        return get_webapp_html(port=self.app.port)

    def initialize_load_config(
        self, request: Union[ChatCompletionRequest, CompletionRequest]
    ) -> LoadConfig:
        """
        Turn the Request object into a partially-complete LoadConfig.

        The load_llm() method is responsible for filling in the rest of
        LoadConfig's parameters.
        """

        # Get model config
        if "/" in request.model:
            # We know the model is a Hugging Face checkpoint if it contains a /
            # This scenario is only supported for run-as-thread
            lc = LoadConfig(model_name="custom", checkpoint=request.model)
        else:
            # The model should be a reference to a built-in model
            lc = LoadConfig(model_name=request.model)

        return lc

    async def completions(
        self, completion_request: CompletionRequest, request: Request
    ):
        """
        Stream completion responses using HTTP chunked transfer encoding.
        """

        lc = self.initialize_load_config(completion_request)

        # Log request parameters (excluding message content for brevity)
        self._log_request_parameters(completion_request, "Completions")

        # Load the model if it's different from the currently loaded one
        await self.load_llm(lc)

        if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
            return self.wrapped_server.completion(completion_request)

        # Check if the model supports reasoning
        reasoning_first_token = self.llm_loaded.reasoning

        # If the model supports reasoning, we:
        # 1. add a <think> tag to the model's context
        # 2. ensure that the first token is a <think> token
        text = completion_request.prompt
        if reasoning_first_token:
            text += "<think>"

        # Prepare generation arguments
        generation_args = {
            "message": text,
            "stop": completion_request.stop,
            "temperature": completion_request.temperature,
            "repeat_penalty": completion_request.repeat_penalty,
            "top_k": completion_request.top_k,
            "top_p": completion_request.top_p,
            "max_new_tokens": completion_request.max_tokens,
        }

        if completion_request.stream:

            if completion_request.logprobs:
                logging.warning("logprobs is not supported for streaming completion")
            if completion_request.echo:
                logging.warning(
                    "`Echo` parameter is not supported for streaming completions"
                )

            # Stream the response
            async def generate():
                # Declare it's the same variable from outside scope
                # This is necessary because the variable is modified
                # in the inner function
                nonlocal reasoning_first_token
                try:
                    async for token in self._generate_tokens(**generation_args):
                        # Handle client disconnect: stop generation and exit
                        if await request.is_disconnected():
                            self.stop_event.set()
                            break

                        choice = CompletionChoice(
                            text=(
                                "<think>" + token if reasoning_first_token else token
                            ),
                            index=0,
                            finish_reason="stop",
                            logprobs=None,
                        )

                        completion = Completion(
                            id="0",
                            choices=[choice],
                            model=self.llm_loaded.checkpoint,
                            object="text_completion",
                            created=int(time.time()),
                        )

                        # Format as SSE
                        reasoning_first_token = False
                        yield f"data: {completion.model_dump_json()}\n\n".encode(
                            "utf-8"
                        )

                    # Send the [DONE] marker only if still connected
                    if not await request.is_disconnected():
                        yield b"data: [DONE]\n\n"
                except asyncio.CancelledError:
                    # Propagate cancellation to the generator loop
                    self.stop_event.set()
                    return

            return StreamingResponse(
                generate(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                },
            )

        # If streaming is not requested, collect all generated tokens into a single response
        else:
            full_response = text if completion_request.echo else ""
            async for token in self._generate_tokens(**generation_args):
                full_response += token

            # If logprobs are requested, create a logprobs object
            logprobs = None
            if completion_request.logprobs:

                # Compute the logprobs
                text_offset, token_logprobs, tokens, top_logprobs = (
                    self.model.compute_logprobs(
                        text=full_response,
                        tokenizer=self.tokenizer,
                        logprobs=completion_request.logprobs,
                    )
                )
                logprobs = Logprobs.model_construct(
                    text_offset=text_offset,
                    token_logprobs=token_logprobs,
                    tokens=tokens,
                    top_logprobs=top_logprobs,
                )

            choice = CompletionChoice(
                text=full_response,
                index=0,
                finish_reason="stop",
                logprobs=logprobs,
            )

            usage = CompletionUsage(
                prompt_tokens=self.input_tokens,
                completion_tokens=self.output_tokens,
                total_tokens=self.input_tokens + self.output_tokens,
            )

            return Completion(
                id="0",
                choices=[choice],
                usage=usage,
                model=self.llm_loaded.checkpoint,
                object="text_completion",
                created=int(time.time()),
            )

    async def chat_completions(
        self, chat_completion_request: ChatCompletionRequest, request: Request
    ):
        """
        Stream chat completion responses using HTTP chunked transfer encoding.
        """

        if chat_completion_request.logprobs:
            logging.warning("logprobs is not supported on chat completion")

        lc = self.initialize_load_config(chat_completion_request)

        # Log request parameters (excluding message history for brevity)
        self._log_request_parameters(chat_completion_request, "Chat completions")

        # Load the model if it's different from the currently loaded one
        await self.load_llm(lc)

        if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
            if (
                hasattr(chat_completion_request, "enable_thinking")
                and chat_completion_request.enable_thinking is False
                and "qwen3" in self.llm_loaded.model_name.lower()
            ):

                # Modify the last user message to include /no_think
                if chat_completion_request.messages:
                    for i in range(len(chat_completion_request.messages) - 1, -1, -1):
                        if chat_completion_request.messages[i].get("role") == "user":
                            original_content = chat_completion_request.messages[i][
                                "content"
                            ]
                            chat_completion_request.messages[i][
                                "content"
                            ] = f"/no_think\n{original_content}"
                            break
            return self.wrapped_server.chat_completion(chat_completion_request)

        # Convert chat messages to text using the model's chat template
        text = self.apply_chat_template(
            chat_completion_request.messages,
            tools=chat_completion_request.tools,
        )

        # If the model supports reasoning, we:
        # 1. add a <think> tag to the model's context
        # 2. ensure that the first token is a <think> token
        reasoning_first_token = self.llm_loaded.reasoning

        if reasoning_first_token:
            text += "<think>"
        # Set the max_new_tokens parameter
        if (
            chat_completion_request.max_completion_tokens
            and chat_completion_request.max_tokens
        ):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=(
                    "Both max_tokens and max_completion_tokens were provided. "
                    "Please use only one of these parameters.",
                ),
            )
        max_new_tokens = (
            chat_completion_request.max_completion_tokens
            if chat_completion_request.max_completion_tokens
            else chat_completion_request.max_tokens
        )

        # Prepare generation arguments
        generation_args = {
            "message": text,
            "stop": chat_completion_request.stop,
            "temperature": chat_completion_request.temperature,
            "repeat_penalty": chat_completion_request.repeat_penalty,
            "top_k": chat_completion_request.top_k,
            "top_p": chat_completion_request.top_p,
            "max_new_tokens": max_new_tokens,
        }

        if chat_completion_request.tools:
            # Get the tool call pattern
            tool_call_pattern = get_tool_call_pattern(
                self.tokenizer.auto_tokenizer.added_tokens_decoder
            )

        if chat_completion_request.stream:

            # Stream the response
            async def generate():
                # Declare it's the same variable from outside scope
                # This is necessary because the variable is modified
                # in the inner function
                nonlocal reasoning_first_token

                # Keep track of the full response for tool call extraction
                full_response = ""

                # Track whether we're still in the thinking phase (before </think> tag)
                in_thinking_phase = self.llm_loaded.reasoning
                reasoning_buffer = ""  # Accumulate reasoning tokens to detect </think>

                try:
                    async for token in self._generate_tokens(**generation_args):
                        # Handle client disconnect: stop generation and exit
                        if await request.is_disconnected():
                            self.stop_event.set()
                            break

                        # Continuously look for tool calls embedded into the generated text
                        openai_tool_calls = None
                        if chat_completion_request.tools:

                            # Append the token to the full response
                            full_response += token

                            tool_calls, _ = extract_tool_calls(
                                full_response,
                                tool_call_pattern,
                            )

                            # If there are tool calls, reset the full response for the next call
                            if tool_calls:
                                openai_tool_calls = []
                                full_response = ""
                            for tool_call in tool_calls:
                                openai_tool_calls.append(
                                    ChoiceDeltaToolCall(
                                        index=0,
                                        id="-",
                                        function=ChoiceDeltaToolCallFunction(
                                            arguments=json.dumps(
                                                tool_call["arguments"]
                                            ),
                                            name=tool_call["name"],
                                        ),
                                        type="function",
                                    )
                                )

                        # Create a ChatCompletionChunk with reasoning_content support
                        # If we're in reasoning mode and haven't seen </think> yet,
                        # send tokens as reasoning_content instead of content
                        delta_content = None
                        delta_reasoning = None

                        if reasoning_first_token:
                            # First token - include opening tag in reasoning
                            delta_reasoning = "<think>" + token
                            reasoning_first_token = False
                            reasoning_buffer = token
                        elif in_thinking_phase:
                            # Still in thinking phase - accumulate and check for </think>
                            reasoning_buffer += token

                            # Check if we've seen the closing tag
                            if "</think>" in reasoning_buffer:
                                # Split at the closing tag
                                before_close, after_close = reasoning_buffer.split(
                                    "</think>", 1
                                )

                                # Send everything before + closing tag as reasoning
                                if before_close or not reasoning_buffer.startswith(
                                    "</think>"
                                ):
                                    delta_reasoning = before_close + "</think>"
                                else:
                                    delta_reasoning = "</think>"

                                # Everything after goes to content (will be sent in next iteration)
                                # For now, mark that we've exited thinking phase
                                in_thinking_phase = False

                                # If there's content after </think>, we need to send it too
                                # But we send it in the current chunk as regular content
                                if after_close:
                                    # We have both reasoning and content in this token
                                    # Send reasoning first, content will accumulate
                                    delta_content = after_close
                            else:
                                # Still accumulating thinking, send as reasoning_content
                                delta_reasoning = token
                        else:
                            # Normal content (after thinking phase ended)
                            delta_content = token

                        chunk = ChatCompletionChunk.model_construct(
                            id="0",
                            object="chat.completion.chunk",
                            created=int(time.time()),
                            model=self.llm_loaded.checkpoint,
                            choices=[
                                Choice.model_construct(
                                    index=0,
                                    delta=ChoiceDelta(
                                        content=delta_content,
                                        reasoning_content=delta_reasoning,
                                        function_call=None,
                                        role="assistant",
                                        tool_calls=openai_tool_calls,
                                        refusal=None,
                                    ),
                                    finish_reason=None,
                                    logprobs=None,
                                )
                            ],
                        )

                        # Format as SSE
                        yield f"data: {chunk.model_dump_json()}\n\n".encode("utf-8")

                    # Send the [DONE] marker only if still connected
                    if not await request.is_disconnected():
                        yield b"data: [DONE]\n\n"
                except asyncio.CancelledError:
                    self.stop_event.set()
                    return

            return StreamingResponse(
                generate(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                },
            )

        # If streaming is not requested, collect all generated tokens into a single response
        else:
            full_response = "<think>" if reasoning_first_token else ""
            async for token in self._generate_tokens(**generation_args):
                full_response += token

            # Extract tool calls from the response
            openai_tool_calls = None
            if chat_completion_request.tools:
                tool_calls, full_response = extract_tool_calls(
                    full_response, tool_call_pattern
                )
                if tool_calls:
                    openai_tool_calls = []
                for tool_call in tool_calls:
                    openai_tool_calls.append(
                        ChatCompletionMessageToolCall(
                            id="-",
                            function=Function(
                                arguments=json.dumps(tool_call["arguments"]),
                                name=tool_call["name"],
                            ),
                            type="function",
                        )
                    )

            ccm = ChatCompletionMessage(
                content=full_response,
                role="assistant",
                refusal=None,
                audio=None,
                function_call=None,
                tool_calls=openai_tool_calls,
            )

            choice = Choice(
                finish_reason="stop",
                index=0,
                message=ccm,
                logprobs=None,
            )

            usage = CompletionUsage(
                prompt_tokens=self.input_tokens,
                completion_tokens=self.output_tokens,
                total_tokens=self.input_tokens + self.output_tokens,
            )

            return ChatCompletion(
                id="0",
                choices=[choice],
                usage=usage,
                model=self.llm_loaded.checkpoint,
                object="chat.completion",
                created=int(time.time()),
            )

    async def embeddings(self, embeddings_request: EmbeddingsRequest):
        """
        Generate embeddings for the provided input.
        """
        # Initialize load config from embeddings request
        lc = LoadConfig(model_name=embeddings_request.model)

        # Load the model if it's different from the currently loaded one
        await self.load_llm(lc)

        if self.llm_loaded.recipe == "llamacpp":
            try:
                return self.wrapped_server.embeddings(embeddings_request)
            except Exception as e:  # pylint: disable=broad-exception-caught
                # Check if model has embeddings label
                model_info = ModelManager().supported_models.get(
                    self.llm_loaded.model_name, {}
                )
                if "embeddings" not in model_info.get("labels", []):
                    raise HTTPException(
                        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                        detail="You tried to generate embeddings for a model that is "
                        "not labeled as an embeddings model. Please use another model "
                        "or re-register the current model with the 'embeddings' label.",
                    ) from e
                else:
                    raise e
        else:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=f"Embeddings not supported for recipe: {self.llm_loaded.recipe}",
            )

    async def reranking(self, reranking_request: RerankingRequest):
        """
        Rerank documents based on their relevance to a query.
        """
        # Initialize load config from reranking request
        lc = LoadConfig(model_name=reranking_request.model)

        # Load the model if it's different from the currently loaded one
        await self.load_llm(lc)

        if self.llm_loaded.recipe == "llamacpp":
            try:
                return self.wrapped_server.reranking(reranking_request)
            except Exception as e:  # pylint: disable=broad-exception-caught
                # Check if model has reranking label
                model_info = ModelManager().supported_models.get(
                    self.llm_loaded.model_name, {}
                )
                if "reranking" not in model_info.get("labels", []):
                    raise HTTPException(
                        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                        detail="You tried to use reranking for a model that is "
                        "not labeled as a reranking model. Please use another model "
                        "or re-register the current model with the 'reranking' label.",
                    ) from e
                else:
                    raise e
        else:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=f"Reranking not supported for recipe: {self.llm_loaded.recipe}",
            )

    def apply_chat_template(
        self, messages: list[dict], tools: list[dict] | None = None
    ):
        """
        Apply the model's chat template to the messages.
        """
        if self.tokenizer.chat_template:

            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                tools=tools,
            )

        # Fallback to a standardized template if the model doesn't provide one
        logging.warning("No chat template found. Using default template.")
        formatted_messages = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            role_marker = "<|assistant|>" if role == "assistant" else "<|user|>"
            formatted_messages.append(f"{role_marker}\n{content} <|end|>")
        return "\n".join(formatted_messages) + "\n<|assistant|>"

    async def responses(self, responses_request: ResponsesRequest, request: Request):
        """
        Stream responses using HTTP chunked transfer encoding.
        """

        lc = self.initialize_load_config(responses_request)

        # Log request parameters (excluding message history for brevity)
        self._log_request_parameters(responses_request, "Responses")

        # Load the model if it's different from the currently loaded one
        await self.load_llm(lc)

        if self.llm_loaded.recipe == "llamacpp":
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=f"Responses API not supported for recipe: {self.llm_loaded.recipe}",
            )

        # Convert chat messages to text using the model's chat template
        if isinstance(responses_request.input, str):
            text = responses_request.input
        else:
            text = self.apply_chat_template(responses_request.input)

        # If the model supports reasoning, we:
        # 1. add a <think> tag to the model's context
        # 2. ensure that the first token is a <think> token
        reasoning_first_token = self.llm_loaded.reasoning

        if reasoning_first_token:
            text += "<think>"

        # Prepare generation arguments
        generation_args = {
            "message": text,
            "temperature": responses_request.temperature,
            "repeat_penalty": responses_request.repeat_penalty,
            "top_k": responses_request.top_k,
            "top_p": responses_request.top_p,
            "max_new_tokens": responses_request.max_output_tokens,
        }

        if responses_request.stream:

            # Stream the response
            async def generate():
                # Declare it's the same variable from outside scope
                # This is necessary because the variable is modified
                # in the inner function
                nonlocal reasoning_first_token

                # Send initial creation event
                response = Response(
                    id="0",
                    model=self.llm_loaded.checkpoint,
                    created_at=int(time.time()),
                    object="response",
                    output=[],
                    parallel_tool_calls=True,
                    tool_choice="auto",
                    tools=[],
                )
                created_event = ResponseCreatedEvent(
                    response=response,
                    type="response.created",
                    sequence_number=0,
                )
                yield f"data: {created_event.model_dump_json()}\n\n".encode("utf-8")

                full_response = "<think>" if reasoning_first_token else ""

                try:
                    async for token in self._generate_tokens(**generation_args):
                        # Handle client disconnect: stop generation and exit
                        if await request.is_disconnected():
                            self.stop_event.set()
                            break

                        # Create an event
                        delta_event = ResponseTextDeltaEvent(
                            content_index=0,
                            delta=(
                                "<think>" + token if reasoning_first_token else token
                            ),
                            item_id="0 ",
                            logprobs=[],
                            output_index=0,
                            sequence_number=0,
                            type="response.output_text.delta",
                        )
                        full_response += token

                        # Format as SSE
                        reasoning_first_token = False
                        yield f"data: {delta_event.model_dump_json()}\n\n".encode(
                            "utf-8"
                        )

                    # Send the completed event (only if still connected)
                    if not await request.is_disconnected():
                        response_output_message = ResponseOutputMessage(
                            id="0",
                            content=[
                                ResponseOutputText(
                                    annotations=[],
                                    text=full_response,
                                    type="output_text",
                                )
                            ],
                            role="assistant",
                            status="completed",
                            type="message",
                        )
                        response = Response(
                            id="0",
                            model=self.llm_loaded.checkpoint,
                            created_at=int(time.time()),
                            object="response",
                            output=[response_output_message],
                            parallel_tool_calls=True,
                            tool_choice="auto",
                            tools=[],
                        )
                        completed_event = ResponseCompletedEvent(
                            response=response,
                            type="response.completed",
                            sequence_number=0,
                        )
                        yield f"data: {completed_event.model_dump_json()}\n\n".encode(
                            "utf-8"
                        )

                        # Send the [DONE] marker
                        yield b"data: [DONE]\n\n"
                except asyncio.CancelledError:
                    self.stop_event.set()
                    return

            return StreamingResponse(
                generate(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                },
            )

        # If streaming is not requested, collect all generated tokens into a single response
        else:
            full_response = "<think>" if reasoning_first_token else ""
            async for token in self._generate_tokens(**generation_args):
                full_response += token

            # Send the completed event
            response_output_message = ResponseOutputMessage(
                id="0",
                content=[
                    ResponseOutputText(
                        annotations=[],
                        text=full_response,
                        type="output_text",
                    )
                ],
                role="assistant",
                status="completed",
                type="message",
            )
            return Response(
                id="0",
                model=self.llm_loaded.checkpoint,
                created_at=int(time.time()),
                object="response",
                output=[response_output_message],
                parallel_tool_calls=True,
                tool_choice="auto",
                tools=[],
            )

    async def _generate_tokens(
        self,
        message: str,
        stop: list[str] | str | None = None,
        max_new_tokens: int | None = None,
        temperature: float | None = None,
        repeat_penalty: float | None = None,
        top_k: int | None = None,
        top_p: float | None = None,
    ):
        """
        Core streaming completion logic, separated from response handling.
        Returns an async generator that yields tokens.
        """

        while not self.app.initialized:
            # Wait for the app's background tasks to finish before
            # allowing generation to proceed
            logging.debug("Waiting for server to fully initialize")
            asyncio.sleep(0.5)
        # These should already be imported as part of the app initialization process,
        # they are just here to make 100% certain and to make the linter happy
        from transformers import TextIteratorStreamer, StoppingCriteriaList

        model = self.model
        tokenizer = self.tokenizer

        # Reset the early-exit flag before we start each generation
        self.stop_event.clear()

        input_ids = tokenizer(message, return_tensors="pt").input_ids

        # Process stop sequences
        stop_sequences = []
        if stop is not None:
            if isinstance(stop, str):
                stop_sequences = [stop]
            else:
                stop_sequences = stop[:4]  # Limit to 4 sequences as per spec

        # Set up the generation parameters
        if "oga-" in self.llm_loaded.recipe:
            from lemonade.tools.oga.utils import OrtGenaiStreamer

            streamer = OrtGenaiStreamer(tokenizer)
            self.input_tokens = len(input_ids)
        else:
            streamer = TextIteratorStreamer(
                tokenizer,
                skip_prompt=True,
            )
            self.input_tokens = len(input_ids[0])

        max_prompt_length = self.ctx_size  # Default fallback
        # For OGA models, try to read the actual max prompt length from config
        if "oga-" in self.llm_loaded.recipe:
            try:
                if model.config and model.config.get("max_prompt_length"):
                    max_prompt_length = model.config["max_prompt_length"]
                    logging.debug(
                        f"Using OGA model max_prompt_length: {max_prompt_length}"
                    )
            # pylint: disable=broad-exception-caught
            except Exception as e:
                logging.debug(f"Could not read OGA model config, using ctx_size: {e}")

        # Apply truncation if input exceeds the limit
        if self.input_tokens > max_prompt_length:
            # Truncate input ids
            truncate_amount = self.input_tokens - max_prompt_length
            input_ids = input_ids[:max_prompt_length]
            # Update token count
            if "oga-" in self.llm_loaded.recipe:
                self.input_tokens = len(input_ids)
            else:
                self.input_tokens = len(input_ids[0])

            # Log warning message instead of raising exception
            truncation_message = (
                f"Input exceeded {max_prompt_length} tokens. "
                f"Truncated {truncate_amount} tokens from the beginning."
            )
            logging.warning(truncation_message)

        # Log the input tokens early to avoid this not showing due to potential crashes
        logging.debug(f"Input Tokens: {self.input_tokens}")
        logging.trace(f"Input Message: {message}")

        if self.llm_loaded.recipe.startswith("hf"):
            stopping_criteria = StoppingCriteriaList([StopOnEvent(self.stop_event)])
        else:
            # HF expects StoppingCriteriaList, which requires torch
            # If we aren't using HF, we can just use a list of StopOnEvent to
            # avoid the torch dep
            stopping_criteria = [StopOnEvent(self.stop_event)]

        generation_kwargs = {
            "input_ids": input_ids,
            "streamer": streamer,
            "max_new_tokens": (
                max_new_tokens if max_new_tokens else DEFAULT_MAX_NEW_TOKENS
            ),
            "min_new_tokens": 1,
            "pad_token_id": tokenizer.eos_token_id,
            "stopping_criteria": stopping_criteria,
            "temperature": temperature,
            "repeat_penalty": repeat_penalty,
            "top_k": top_k,
            "top_p": top_p,
        }

        # Initialize performance variables
        generation_start_time = time.perf_counter()
        first_token = True
        self.decode_token_times = []
        self.output_tokens = 0

        # Begin generation
        thread = GeneratorThread(
            streamer, target=model.generate, kwargs=generation_kwargs
        )
        thread.start()

        # Acquire the generation semaphore
        await self._generate_semaphore.acquire()
        active_generations = (
            self.max_concurrent_generations
            - self._generate_semaphore._value  # pylint: disable=protected-access
        )

        logging.debug(f"Active generations: {active_generations}")

        try:
            # Generate the response using streaming
            new_text = ""
            for new_text in streamer:
                # Yield control back to the event loop
                # This gives the FastAPI server a chance to send the chunks to the client
                await asyncio.sleep(0)

                # Capture performance stats about this token
                self.output_tokens = self.output_tokens + 1
                if first_token:
                    self.time_to_first_token = (
                        time.perf_counter() - generation_start_time
                    )
                    first_token = False
                else:
                    self.decode_token_times.append(
                        time.perf_counter() - next_token_start_time
                    )
                next_token_start_time = time.perf_counter()

                # Remove the EOS token from the response if needed
                if hasattr(self.tokenizer, "eos_token"):
                    new_text = new_text.replace(self.tokenizer.eos_token, "")

                # Check for stop sequences
                if stop_sequences:
                    for stop_seq in stop_sequences:
                        if stop_seq in new_text:
                            # Make sure we yield the text up to before the stop sequence
                            new_text = new_text[: new_text.find(stop_seq)]
                            self.stop_event.set()

                yield new_text

                # Allow the user to finish the response early
                if self.stop_event.is_set():
                    logging.info("Stopping generation early.")
                    break

            if len(self.decode_token_times) > 0:
                self.tokens_per_second = 1 / statistics.mean(self.decode_token_times)
            else:
                self.tokens_per_second = 0

        finally:
            thread.join()

            # Release the semaphore when generation is complete (or if an error occurs)
            self._generate_semaphore.release()
            active_generations = (
                self.max_concurrent_generations
                - self._generate_semaphore._value  # pylint: disable=protected-access
            )

            # Check if an exception occurred in the generation thread
            # If it did, raise it as an HTTPException so that the client
            # knows they wont be getting a completion
            if thread.exception:
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=f"Completion failure: {thread.exception}",
                )

            # Display telemetry if in debug mode
            await self._show_telemetry()

    async def send_stats(self):
        """
        Send performance statistics to the client.
        """
        # If using wrapped server, get telemetry from the telemetry instance
        if self.llm_loaded and (
            self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm"
        ):
            return self.wrapped_server.telemetry.get_telemetry_data()

        # For built-in server, use the existing telemetry
        return {
            "time_to_first_token": self.time_to_first_token,
            "tokens_per_second": self.tokens_per_second,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "decode_token_times": self.decode_token_times,
        }

    async def halt_generation(self):
        """
        Allow the client to halt an in-progress generation.
        """

        self.stop_event.set()

        return {
            "terminated": True,
        }

    async def health(self):
        """
        Report server health information to the client.
        """

        return {
            "status": "ok",
            "checkpoint_loaded": (
                self.llm_loaded.checkpoint if self.llm_loaded else None
            ),
            "model_loaded": (
                self.llm_loaded.model_name
                if (self.llm_loaded and self.llm_loaded.model_name)
                else None
            ),
        }

    async def get_system_info(self, request: Request):
        """
        Return system and device enumeration information.
        Supports optional 'verbose' query parameter.
        """
        from lemonade.common.system_info import (
            get_system_info_dict,
            get_device_info_dict,
            get_system_info as get_system_info_obj,
        )

        # Get verbose parameter from query string (default to False)
        verbose = request.query_params.get("verbose", "false").lower() in ["true", "1"]

        info = get_system_info_dict()
        info["devices"] = get_device_info_dict()

        # Filter out verbose-only information if not in verbose mode
        if not verbose:
            essential_keys = ["OS Version", "Processor", "Physical Memory", "devices"]
            info = {k: v for k, v in info.items() if k in essential_keys}
        else:
            # In verbose mode, add Python packages at the end
            system_info_obj = get_system_info_obj()
            info["Python Packages"] = system_info_obj.get_python_packages()

        return info

    def model_load_failure(self, model_reference: str, message: Optional[str] = None):
        """
        Clean up after a model load failure, then log it and raise
        an HTTPException with details.
        """
        self.llm_loaded = None
        self.tokenizer = None
        self.model = None

        default_message = "see stack trace and error message below"
        if message:
            detail = message
        else:
            detail = default_message

        logging.exception(f"Tried to load LLM {model_reference} and failed: {detail}")

        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=detail,
        )

    async def pull(self, config: PullConfig):
        """
        Install a supported LLM by its Lemonade Model Name.
        """

        # Install the model
        ModelManager().download_models(
            [config.model_name],
            checkpoint=config.checkpoint,
            recipe=config.recipe,
            reasoning=config.reasoning,
            vision=config.vision,
            mmproj=config.mmproj,
            # The pull endpoint will download an upgraded model if available, even
            # if we already have a local copy of the model
            do_not_upgrade=False,
        )

        # Refresh the list of downloaded models, to ensure it
        # includes the model we just installed
        self.local_models = ModelManager().downloaded_models_enabled

    async def delete(self, config: DeleteConfig):
        """
        Delete a supported LLM by its Lemonade Model Name.
        """
        try:
            # If the model to be deleted is currently loaded, unload it first
            if self.llm_loaded and self.llm_loaded.model_name == config.model_name:
                await self.unload_llm(require_lock=True)

            # Delete the model
            ModelManager().delete_model(config.model_name)

            # Refresh the list of downloaded models
            self.local_models = ModelManager().downloaded_models_enabled

            return {
                "status": "success",
                "message": f"Deleted model: {config.model_name}",
            }
        except ValueError as e:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=str(e),
            )
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Failed to delete model {config.model_name}: {str(e)}",
            )

    async def load_llm(self, config: LoadConfig):
        """
        Load a registered LLM into system memory. Install the model first, if needed.
            config: the information required to load the model
        """
        from huggingface_hub.constants import HF_HUB_CACHE

        try:
            await self._load_lock.acquire()
            # Acquire all generate locks
            for _ in range(self.max_concurrent_generations):
                await self._generate_semaphore.acquire()

            # Make sure the model is already registered
            supported_models = ModelManager().supported_models

            # The `custom` name allows run-as-thread servers to bypass loading
            if config.model_name == "custom":
                config_to_use = config
            else:
                if config.model_name not in supported_models.keys():
                    self.model_load_failure(
                        config.model_name,
                        message=(
                            f"Load request for model_name={config.model_name} "
                            "not registered with Lemonade Server. You can register and "
                            "install new models with a `pull` request."
                        ),
                    )

                # Get additional properties from the model registry
                config_to_use = LoadConfig(**supported_models[config.model_name])

            # For locally uploaded models, convert the relative checkpoint path to absolute path
            model_source = supported_models.get(config.model_name, {}).get(
                "source", None
            )
            if (
                model_source == "local_upload"
                and config_to_use.checkpoint
                and not config_to_use.recipe.startswith("hf-")
            ):
                # Check if checkpoint is a relative path (stored during upload)
                if not os.path.isabs(config_to_use.checkpoint):
                    # Convert relative path to absolute by joining with HF_HUB_CACHE
                    absolute_checkpoint = os.path.join(
                        HF_HUB_CACHE, config_to_use.checkpoint
                    )
                    if os.path.exists(absolute_checkpoint):
                        config_to_use.checkpoint = absolute_checkpoint
                    else:
                        logging.warning(
                            f"Checkpoint path does not exist: {absolute_checkpoint}"
                        )

                # Also resolve mmproj path if present
                if config_to_use.mmproj and not os.path.isabs(config_to_use.mmproj):
                    absolute_mmproj = os.path.join(HF_HUB_CACHE, config_to_use.mmproj)
                    if os.path.exists(absolute_mmproj):
                        config_to_use.mmproj = absolute_mmproj
                    else:
                        logging.warning(
                            f"MMProj path does not exist: {absolute_mmproj}"
                        )

            # Caching mechanism: if the checkpoint is already loaded there is nothing else to do
            if (
                self.llm_loaded
                and config_to_use.checkpoint == self.llm_loaded.checkpoint
            ):
                if (
                    self.llm_loaded.recipe == "llamacpp"
                    or self.llm_loaded.recipe == "flm"
                ) and self.wrapped_server.process.poll():
                    # wrapped server process has gone away for some reason, so we should
                    # proceed with loading to get it back
                    pass
                else:
                    return {
                        "status": "success",
                        "message": f"Model already loaded: {config.model_name}",
                    }

            # Unload the current model if needed
            if self.llm_loaded:
                await self.unload_llm(require_lock=False)

            logging.info(f"Loading llm: {config.model_name}")
            try:
                if config_to_use.recipe == "llamacpp":
                    self.wrapped_server = LlamaServer(self.llamacpp_backend)
                    self.wrapped_server.load(
                        model_config=config_to_use,
                        ctx_size=self.ctx_size,
                        do_not_upgrade=True,
                    )

                elif config_to_use.recipe == "flm":
                    self.wrapped_server = FlmServer()
                    self.wrapped_server.load(
                        model_config=config_to_use,
                        ctx_size=self.ctx_size,
                        do_not_upgrade=True,
                    )

                else:
                    self.model, self.tokenizer = lemonade_api.from_pretrained(
                        checkpoint=config_to_use.checkpoint, recipe=config_to_use.recipe
                    )
                self.llm_loaded = config_to_use

                return {
                    "status": "success",
                    "message": f"Loaded model: {config.model_name}",
                }
            except HTTPException:
                raise
            except Exception:  # pylint: disable=broad-exception-caught
                self.model_load_failure(config.model_name)

        finally:
            self._load_lock.release()

            # Release all generate locks
            for _ in range(self.max_concurrent_generations):
                self._generate_semaphore.release()

            # Refresh the list of downloaded models, to ensure it
            # includes the model we just loaded
            if config.model_name not in self.local_models:
                self.local_models = ModelManager().downloaded_models_enabled

    async def unload_llm(self, require_lock: bool = True):
        try:
            if require_lock:
                await self._load_lock.acquire()

                # Acquire all generate locks
                for _ in range(self.max_concurrent_generations):
                    await self._generate_semaphore.acquire()

            if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
                self.wrapped_server.process.terminate()

            self.llm_loaded = None
            self.tokenizer = None
            self.model = None
            return {"status": "success", "message": "Unloaded model"}
        except Exception as e:  # pylint: disable=broad-exception-caught
            return {
                "status": "error",
                "message": f"Failed to unload model: {str(e)}",
            }
        finally:
            if require_lock:
                self._load_lock.release()

                # Release all generate locks
                for _ in range(self.max_concurrent_generations):
                    self._generate_semaphore.release()

    async def models(self):
        """
        Return a list of available models in OpenAI-compatible format.
        """
        models_list = []
        for model in self.local_models:
            m = ServerModel(
                id=model,
                owned_by="lemonade",
                object="model",
                created=int(time.time()),
                checkpoint=self.local_models[model]["checkpoint"],
                recipe=self.local_models[model]["recipe"],
            )
            models_list.append(m)

        return {"object": "list", "data": models_list}

    async def retrieve_model(self, model_id: str):
        """
        Retrieve a specific model by ID in OpenAI-compatible format.
        """
        # Raise an error if the model does not exist
        if model_id not in self.local_models:
            # Mimic the error format of the OpenAI API
            raise HTTPException(
                status_code=404,
                detail={
                    "message": f"model {model_id} not found",
                    "type": "api_error",
                    "param": None,
                    "code": None,
                },
            )

        # Return the specific model
        model_info = self.local_models[model_id]
        model = ServerModel(
            id=model_id,
            owned_by="lemonade",
            object="model",
            created=int(time.time()),
            checkpoint=model_info["checkpoint"],
            recipe=model_info["recipe"],
        )

        return model

    def setup_middleware_timer(self):
        logging.info("Middleware set up")

        @self.app.middleware("http")
        async def log_request_time(request: Request, call_next):
            """
            Log the request processing time for any request.
            For streaming responses, wraps the body iterator to measure total time.
            Only applies the wrapper in debug mode.
            """
            start_time = time.perf_counter()
            response = await call_next(request)

            if (
                self.debug_logging_enabled
                and hasattr(response, "body_iterator")
                and response.body_iterator is not None
            ):
                original_iterator = response.body_iterator

                async def wrapped_iterator():
                    async for chunk in original_iterator:
                        yield chunk
                    request_time = time.perf_counter() - start_time
                    logging.debug(
                        f"Total request time (streamed): {request_time:.4f} seconds"
                    )

                response.body_iterator = wrapped_iterator()
            else:
                request_time = time.perf_counter() - start_time
                if self.debug_logging_enabled:
                    logging.debug(f"Total request time: {request_time:.4f} seconds")
            return response

    async def logs_ws(self, websocket: WebSocket):
        if not self.log_file or not os.path.exists(self.log_file):
            await websocket.close(code=4000)
            return
        await log_streamer(websocket, self.log_file)

    async def get_incompatible_models(self):
        """
        Get information about incompatible RyzenAI models in the cache.
        """
        try:
            return ModelManager().get_incompatible_ryzenai_models()
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Failed to scan for incompatible models: {str(e)}",
            )

    async def cleanup_incompatible_models(self, request: Request):
        """
        Delete selected incompatible RyzenAI models from the cache.
        """
        try:
            body = await request.json()
            model_paths = body.get("model_paths", [])

            if not model_paths:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail="No model_paths provided",
                )

            result = ModelManager().cleanup_incompatible_models(model_paths)
            return result
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Failed to cleanup models: {str(e)}",
            )


# This file was originally licensed under Apache 2.0. It has been modified.
# Modifications Copyright (c) 2025 AMD
