import os
import gradio as gr
from gradio import ChatMessage
from typing import Iterator
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

# Загрузка модели и токенизатора
model_name = "FractalGPT/RuQwen2.5-3B-Instruct-AWQ"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

# Создание пайплайна для генерации текста
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def format_chat_history(messages: list) -> str:
    """
    Форматирует историю чата в строку, которую модель может понять.
    """
    formatted_history = ""
    for message in messages:
        if message.get("role") == "user":
            formatted_history += f"User: {message.get('content', '')}\n"
        elif message.get("role") == "assistant":
            formatted_history += f"Assistant: {message.get('content', '')}\n"
    return formatted_history

def stream_model_response(user_message: str, messages: list) -> Iterator[list]:
    """
    Генерирует ответ модели с поддержкой истории чата.
    """
    try:
        print(f"\n=== New Request ===")
        print(f"User message: {user_message}")
        
        # Форматируем историю чата
        chat_history = format_chat_history(messages)
        
        # Формируем входной текст для модели
        input_text = f"{chat_history}User: {user_message}\nAssistant:"
        
        # Генерируем ответ модели
        response = text_generator(input_text, max_length=512, do_sample=True, temperature=0.7, top_p=0.9)
        model_response = response[0]['generated_text'].split("Assistant:")[-1].strip()
        
        # Добавляем ответ модели в историю чата
        messages.append(
            ChatMessage(
                role="assistant",
                content=model_response
            )
        )
        
        yield messages
        
        print(f"\n=== Final Response ===\n{model_response}")
                
    except Exception as e:
        print(f"\n=== Error ===\n{str(e)}")
        messages.append(
            ChatMessage(
                role="assistant",
                content=f"I apologize, but I encountered an error: {str(e)}"
            )
        )
        yield messages

def user_message(msg: str, history: list) -> tuple[str, list]:
    """Добавляет сообщение пользователя в историю чата"""
    history.append(ChatMessage(role="user", content=msg))
    return "", history
    

# Создаем интерфейс Gradio
with gr.Blocks(theme=gr.themes.Citrus(), fill_height=True) as demo:
    gr.Markdown("# Chat with FractalGPT/RuQwen2.5-3B-Instruct-AWQ 💭")

    chatbot = gr.Chatbot(
        type="messages",
        label="FractalGPT Chatbot",
        render_markdown=True,
        scale=1,
        avatar_images=(None, "https://huggingface.co/FractalGPT/RuQwen2.5-3B-Instruct-AWQ/resolve/main/avatar.png")
    )

    with gr.Row(equal_height=True):
        input_box = gr.Textbox(
            lines=1,
            label="Chat Message",
            placeholder="Type your message here...",
            scale=4
        )

        clear_button = gr.Button("Clear Chat", scale=1)

    # Настраиваем обработчики событий
    msg_store = gr.State("")  # Хранилище для сохранения сообщения пользователя
    
    input_box.submit(
        lambda msg: (msg, msg, ""),  # Сохраняем сообщение и очищаем поле ввода
        inputs=[input_box],
        outputs=[msg_store, input_box, input_box],
        queue=False
    ).then(
        user_message,  # Добавляем сообщение пользователя в чат
        inputs=[msg_store, chatbot],
        outputs=[input_box, chatbot],
        queue=False
    ).then(
        stream_model_response,  # Генерируем и передаем ответ модели
        inputs=[msg_store, chatbot],
        outputs=chatbot
    )

    clear_button.click(
        lambda: ([], "", ""),
        outputs=[chatbot, input_box, msg_store],
        queue=False
    )

# Запускаем интерфейс
if __name__ == "__main__":
    demo.launch(debug=True)