from typing import Iterable
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from openai import OpenAI, Stream
from openai.types.chat import (
    ChatCompletionMessageParam,
    ChatCompletionChunk,
    ChatCompletion,
)
from openai.types.responses import ToolParam
from pydantic import BaseModel

app = FastAPI()


class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[ChatCompletionMessageParam]
    tools: list[ToolParam] = []
    stream: bool = False


@app.get("/v1/models")
async def get_models():
    return {
        "object": "list",
        "data": [{"id": "gpt-4o-mini", "object": "model", "owned_by": "openai"}],
    }


def to_sse_stream(models: Stream[ChatCompletionChunk]) -> Iterable[bytes]:
    for model in models:
        chunk = f"data: {model.model_dump_json(ensure_ascii=False)}\n\n"
        yield chunk.encode("utf-8")

    yield "data: [DONE]".encode("utf-8")


@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    client = OpenAI()

    completion = client.chat.completions.create(
        model=request.model,
        messages=request.messages,
        stream=request.stream,
    )

    if isinstance(completion, ChatCompletion):
        return completion.model_dump()
    return StreamingResponse(to_sse_stream(completion), media_type="text/event-stream")
