from fastapi import Body, FastAPI
from ollama import ChatResponse, Client

from config import OLLAMA_MODEL, OLLAMA_URL

client = Client(host=f"{OLLAMA_URL}")
# ollama.create(model='example', from_='gemma3', system="You are Mario from Super Mario Bros.") // TODO: Для асистента?


def message(text: str) -> ChatResponse:
    return chat(
        [
            {
                "role": "system",
                "content": "Отвечай строго в формате Markdown. Не нужно пихать везде большие заголовки! Пиши как обычный человек, но красиво оформляй ответ",
            },
            {"role": "user", "content": text},
        ]
    )


def chat(messages: list) -> ChatResponse:
    return client.chat(model=f"{OLLAMA_MODEL}", messages=messages, think=True)


app = FastAPI()


@app.post("/chat")
def _(prompt: str = Body(..., embed=True)):
    print(prompt)
    response = message(prompt)

    return response


@app.post("/generate")
def generate(prompt: str = Body(..., embed=True)):
    response = client.generate(
        model=f"{OLLAMA_MODEL}",
        prompt=prompt,
        think=True,
        options={"temperature": 0.15},
    )

    return response


@app.get("/health")
def health():
    return {"status": "ok"}