import logging
from fastapi import FastAPI
from llama_cpp import Llama

app = FastAPI()

CHAT_TEMPLATE = '<|system|> {system_prompt}<|end|><|user|> {prompt}<|end|><|assistant|>'.strip()

logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

logger.info("Запускаемся... 🥳🥳🥳")

# Инициализация модели
try:
    # загрузка модели для локального хранилища
    # llm = Llama(
    #     model_path="./models/phi-3-mini-4k-instruct-q4.gguf",
    #     verbose=False,
    #     n_gpu_layers=-1,
    #     n_ctx=4096
    # )

    logger.info("Загрузка модели...")
    llm = Llama.from_pretrained(
        repo_id='bartowski/Phi-3.5-mini-instruct-GGUF',
        filename='Phi-3.5-mini-instruct-Q6_K_L.gguf',
        n_gpu_layers=-1,
        n_ctx=4096,
    )

except Exception as e:
    logger.error(f"Ошибка загрузки модели: {str(e)}")
    raise


# составление промта для модели
def create_prompt(text: str) -> str | None:
    try:
        user_input = text
        logger.info(f"Получено сообщение: {user_input}")

        system_prompt = 'Ответ должен быть точным и кратким и если возможно шутливым.'

        # Генерация шаблона
        return CHAT_TEMPLATE.format(
            system_prompt=system_prompt,
            prompt=user_input,
        )
    except Exception as e:
        logger.error(e)


def generate_response(prompt: str) -> str:
    try:
        # Обработка текстового сообщения
        output = llm(
            prompt,
            max_tokens=512,
            stop=["<|end|>"],
        )

        logger.info('Output:')
        logger.info(output)

        response: str = output['choices'][0]['text']

        # Отправка ответа
        if response:
            return response

        return 'Произошла ошибка при обработке запроса'
    except Exception as e:
        logger.error(f"Ошибка обработки сообщения: {str(e)}")


@app.post("/predict")
async def predict(text: str):
    # Генерация ответа с помощью модели
    prompt = create_prompt(text)
    response = generate_response(prompt)
    return {"response": response}