# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium

import threading
import time
import uuid

import torch


try:
    import uvicorn
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel
except ModuleNotFoundError as exception:
    raise type(exception)(
        "GPT-QModel OpenAi serve required dependencies are not installed.",
        "Please install via `pip install gptqmodel[openai] --no-build-isolation`.",
    )

class OpenAiServer:
    def __init__(self, model):
        self.uvicorn_server = None
        self.app = FastAPI()
        self.model = model
        self.tokenizer = model.tokenizer
        self.model_id_or_path = model.config.name_or_path
        self.setup_routes()

    def setup_routes(self):
        class OpenAiRequest(BaseModel):
            model: str
            messages: list = []
            max_tokens: int = 256
            temperature: float = 0.0
            top_p: float = 1.0
            n: int = 1
            stop: list = None

        class OpenAiResponseChoice(BaseModel):
            text: str
            index: int = 0

        class OpenAiResponse(BaseModel):
            id: str = ""
            object: str = "text_completion"
            created: int
            model: str
            choices: list[OpenAiResponseChoice]

        @self.app.post("/v1/chat/completions", response_model=OpenAiResponse)
        async def create_completion(request: OpenAiRequest):
            try:
                model_inputs = self.tokenizer.apply_chat_template(
                    request.messages,
                    add_generation_prompt=True,
                    return_tensors='pt',
                )

                if isinstance(model_inputs, torch.Tensor):
                    model_inputs = model_inputs.to(self.model.device)
                    generate_inputs = {"inputs": model_inputs}
                    prompt_length = model_inputs.size(-1)
                else:
                    model_inputs = model_inputs.to(self.model.device)
                    input_ids = model_inputs["input_ids"]
                    generate_inputs = dict(model_inputs)
                    prompt_length = input_ids.size(-1)

                do_sample = True if request.temperature != 0.0 else False
                with torch.inference_mode():
                    outputs = self.model.generate(
                        max_length=prompt_length + request.max_tokens,
                        temperature=request.temperature,
                        top_p=request.top_p,
                        num_return_sequences=request.n,
                        eos_token_id=self.tokenizer.eos_token_id,
                        stop_strings=request.stop,
                        do_sample=do_sample,
                        **generate_inputs,
                    )

                generated_texts = self.tokenizer.batch_decode(
                    outputs[:, prompt_length:],
                    skip_special_tokens=True,
                )

                choices = [
                    OpenAiResponseChoice(
                        text=gen_text,
                        index=i,
                    )
                    for i, gen_text in enumerate(generated_texts)
                ]

                response = OpenAiResponse(
                    id=f"{uuid.uuid4()}",
                    created=int(time.time()),
                    model=self.model_id_or_path,
                    choices=choices
                )
                return response
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.get("/")
        def read_root():
            return {"message": "GPT-QModel OpenAI Compatible Server is running."}

        @self.app.get("/shutdown")
        def shutdown():
            self.shutdown()
            return {"message": "Server is shutting down..."}

    def start(self, host: str = "0.0.0.0", port: int = 80, async_mode: bool = True):
        config = uvicorn.Config(self.app, host=host, port=port, log_level="info")
        self.uvicorn_server = uvicorn.Server(config)

        def run_server():
            self.uvicorn_server.run()

        if async_mode:
            thread = threading.Thread(target=run_server, daemon=False)
            thread.start()
            print(f"GPT-QModel OpenAi Server has started asynchronously at http://{host}:{port}.")
        else:
            run_server()
            print(f"GPT-QModel OpenAi Server has started synchronously at http://{host}:{port}.")

    def shutdown(self):
        if self.uvicorn_server is not None:
            self.uvicorn_server.should_exit = True
            print("GPT-QModel OpenAi Server is shutting down...")

    def wait_until_ready(self, timeout: int = 30, check_interval: float = 0.1):
        start_time = time.time()
        while not self.uvicorn_server.started:
            if time.time() - start_time > timeout:
                raise TimeoutError("GPT-QModel OpenAi server failed to start within the specified time.")
            time.sleep(check_interval)