# imports
import gradio as gr
import requests
import json
import os


# functions
def generate(description):
    if not description:
        yield None
        return

    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {os.getenv("API_KEY")}'
    }

    payload = {
        'messages': [{'role': 'system', 'content': f'Ты - корректор текста. Пользователь будет отправлять тебе сообщения, а ты должен исправлять в них все ГРАММАТИЧЕСКИЕ ошибки (в том числе и знаки припенания). Отвечать ты должен без лишних символов, только исправленный запрос пользователя, без markdown. Не надо писать вроде этого: "Вот ваш исправленный тексст: ...", пиши сразу "...". И не нужно выделять то, что ты исправил, просто напиши всё правильно, без грамматических ошибок. Начинаем, ниже сообщение пользователя'}, {'role': 'user', 'content': description}],
        'max_tokens': 150000,
        'model': "gemini-1.5-pro-latest",
        'stream': True # Добавляем stream=True
    }

    try:
        response = requests.post(f'{os.getenv("BASE_URL")}v1/chat/completions', headers=headers, json=payload, stream=True, timeout=200)
        response.raise_for_status()  # Проверяем на ошибки HTTP

        full_text = ""
        for chunk in response.iter_lines():
            if chunk:
                try:
                    chunk = chunk.decode('utf-8').replace("data: ", "")
                    if chunk == "[DONE]":
                        break
                    chunk_data = json.loads(chunk)
                    if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                        text_chunk = chunk_data['choices'][0]['delta'].get('content', "")
                        full_text += text_chunk
                        yield full_text  # Выводим текст как итератор
                except json.JSONDecodeError:
                    continue
        if not full_text:
             yield "Не удалось получить ответ от сервера."
    except requests.exceptions.RequestException as e:
        print(f"Ошибка запроса: {e}")
        yield f"**Ошибка запроса!**\n\n```\n{e}\n```"
    except Exception as e:
        print(f"Ошибка: {str(e)}")
        yield "Произошла ошибка при генерации"


# Ссылка на файл CSS
css_url = "https://neurixyufi-aihub.static.hf.space/style.css"

# Получение CSS по ссылке
response = requests.get(css_url)
css = response.text + " .gradio-container{max-width: 700px !important} h1{text-align:center}"


# ui
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Исправление опечаток")
    with gr.Tab("Исправление опечаток"):
        with gr.Row():
            promt = gr.Textbox(show_label=True, label="Запрос", lines=3, placeholder="Привет как делпа ?")
        with gr.Row():
            text_button = gr.Button("Генерация", variant='primary')
    with gr.Row():
        with gr.Tab("Ответ"):
            text_output = gr.Textbox(show_label=False, placeholder="Привет, как дела?")

    text_button.click(generate, inputs=[promt], outputs=[text_output], concurrency_limit=250)

demo.queue(api_open=False).launch()