from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import torch

class MultiPersonalityChatBot:
    def __init__(self):
        self.personalities = {
            "Гаргамель": {
                "model": "sberbank-ai/rugpt3large_based_on_gpt2",
                "prompt": "Гаргамель: Я великий волшебник! ",
                "temperature": 0.9
            },
            "Родя": {
                "model": "sberbank-ai/rugpt3large_based_on_gpt2",
                "prompt": "Родя: Привет, братан! ",
                "temperature": 0.8
            },
            "Ромчик": {
                "model": "sberbank-ai/rugpt3large_based_on_gpt2", 
                "prompt": "Ромчик: Слушай, вот что я думаю... ",
                "temperature": 0.7
            },
            "Кузя": {
                "model": "RussianSuperGLUE/rugpt3large_based_on_gpt2",
                "prompt": "Кузя: Опа, народ! ",
                "temperature": 0.85
            },
            "Nekale": {
                "model": "sberbank-ai/rugpt3large_based_on_gpt2",
                "prompt": "Nekale: Йоу, всем салют! ",
                "temperature": 0.75
            },
            "Ваня Парилка": {
                "model": "RussianSuperGLUE/rugpt3large_based_on_gpt2",
                "prompt": "Ваня Парилка: Знаток бани и здоровья! ",
                "temperature": 0.6
            }
        }
        self.current_personality = "Nekale"
        self.tokenizer = None
        self.model = None
        self.load_model(self.current_personality)
    
    def load_model(self, personality):
        if personality in self.personalities:
            model_name = self.personalities[personality]["model"]
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForCausalLM.from_pretrained(model_name)
            self.current_personality = personality
            return f"✅ Загружена личность: {personality}"
        return "❌ Ошибка загрузки личности"
    
    def generate_response(self, message, history):
        if not self.model or not self.tokenizer:
            return "❌ Сначала загрузите личность"
            
        personality_data = self.personalities[self.current_personality]
        prompt = f"{personality_data['prompt']}Вопрос: {message}\nОтвет:"
        
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs.input_ids,
                max_length=inputs.input_ids.shape[1] + 150,
                num_return_sequences=1,
                temperature=personality_data['temperature'],
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                repetition_penalty=1.1,
                no_repeat_ngram_size=3
            )
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response.replace(prompt, "").strip()
        return response

chatbot = MultiPersonalityChatBot()

with gr.Blocks(title="Мульти-Персональный Чат-Бот", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎭 Мульти-Персональный Чат-Бот")
    gr.Markdown("## Выберите персонажа и общайтесь!")
    
    with gr.Row():
        with gr.Column(scale=1):
            personality = gr.Dropdown(
                choices=list(chatbot.personalities.keys()),
                label="Выберите персонажа",
                value="Nekale"
            )
            load_btn = gr.Button("🎯 Активировать персонажа", variant="primary")
            status = gr.Textbox(label="Статус системы", interactive=False)
            
            gr.Markdown("### 🎪 Наши персонажи:")
            gr.Markdown("""
            - **Гаргамель** - Великий волшебник
            - **Родя** - Братан для душевных бесед  
            - **Ромчик** - Мудрый собеседник
            - **Кузя** - Душа компании
            - **Nekale** - Современный и стильный
            - **Ваня Парилка** - Эксперт по здоровью и бане
            """)
        
        with gr.Column(scale=2):
            chat_interface = gr.ChatInterface(
                fn=chatbot.generate_response,
                title="💬 Чат с персонажем",
                description="Начните общение с выбранным персонажем...",
                examples=[
                    "Расскажи о себе",
                    "Что ты умеешь?",
                    "Дай совет по здоровью"
                ]
            )

    load_btn.click(
        fn=chatbot.load_model,
        inputs=personality,
        outputs=status
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )