# app.py
import gradio as gr
from g4f.client import Client
import base64
from PIL import Image
import io

client = Client()

# --- Функции для очистки ответов ---
import re
from typing import List, Tuple, Any

def strip_md_refs(text: str) -> str:
    """
    Удаляет из строки все конструкции вида [[N]](URL),
    где N – произвольное число (может быть и несколько цифр).
    Возвращает «чистый» текст без ссылок.
    """
    pattern = r'\[\[\d+\]\]\([^\)]*\)'
    return re.sub(pattern, '', text)

def strip_citations(text: str) -> str:
    """
    Удаляет из текста все строки вида:
    а также любые блоки, начинающиеся с символа '>' (цитаты).
    Возвращает «чистый» текст без «концовки».
    """
    text = re.sub(r'^\s*>.*(?:\n|$)', '', text, flags=re.MULTILINE)
    text = re.sub(r'^\s*\[\d+\]\s*.+(?:\n|$)', '', text, flags=re.MULTILINE)
    text = re.sub(r'\n{2,}', '\n\n', text).strip()
    return text

def image_to_base64(image):
    """Конвертирует изображение в base64 строку"""
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

# --- Обработчик диалога ---
def respond(message, history, system_prompt, image=None):
    # Используем пользовательский промт из интерфейса
    messages = [{"role": "system", "content": system_prompt}]
    
    # Добавляем историю диалога
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    
    # Формируем текущее сообщение
    current_message = {"role": "user", "content": message}
    
    # Если есть изображение, добавляем его к сообщению
    if image is not None:
        image_base64 = image_to_base64(image)
        current_message["images"] = [image_base64]
    
    messages.append(current_message)

    try:
        responsed = client.chat.completions.create(
            model="gpt-4",
            messages=messages,
            temperature=0.8,
            top_p=0.95,
            max_tokens=581691,
            presence_penalty=0.2,
            frequency_penalty=0.1,
        )
        bot_message = strip_citations(f"{strip_md_refs(responsed.choices[0].message.content)}")

    except Exception as e:
        bot_message = f"Ошибка: {str(e)}"
    return history + [[message, bot_message]]

# --- Интерфейс ---
with gr.Blocks(title="ESP Brain") as demo:
    gr.Markdown("## ESP Brain - Настраиваемый ассистент")

   # Поле для ввода системного промта
    system_prompt_input = gr.Textbox(
    label="Системный промт",
    value=""" """,
    lines=5,
    max_lines=10,
    placeholder="Введите системный промт для настройки поведения ассистента..."
)

    chatbot = gr.Chatbot(
    height=600,
)

    with gr.Row():
        txt = gr.Textbox(
        placeholder="Напиши сообщение...",
        show_label=False,
        scale=8
    )
        submit_btn = gr.Button("Отправить", scale=2)

    with gr.Row():
        retry_btn = gr.Button("🔄 Повторить")
        undo_btn = gr.Button("↩️ Отменить")
        clear_btn = gr.Button("🗑️ Очистить")

# Логика
    def submit_message(message, history, system_prompt):
        return respond(message, history, system_prompt)

    txt.submit(fn=submit_message, inputs=[txt, chatbot, system_prompt_input], outputs=chatbot)
    submit_btn.click(fn=submit_message, inputs=[txt, chatbot, system_prompt_input], outputs=chatbot)


    def retry_last(history, system_prompt):
        if history:
            last_user_msg = history[-1][0]
            new_history = history[:-1]
            return respond(last_user_msg, new_history, system_prompt, None)
        return history

    def undo_last(history):
        return history[:-1]

    retry_btn.click(fn=retry_last, inputs=[chatbot, system_prompt_input], outputs=chatbot, queue=False)
    undo_btn.click(fn=undo_last, inputs=chatbot, outputs=chatbot, queue=False)
    clear_btn.click(lambda: [], outputs=chatbot, queue=False)

# --- Запуск ---
if __name__ == "__main__":
    demo.queue()
    demo.launch(
        share=True,
        ssr_mode=False,
        debug=True
    )