import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json
from pathlib import Path
from datetime import datetime
import requests
import re

# === ЗАГРУЗКА МОДЕЛИ ===
print("🚀 Загрузка модели...")
model_id = "OpenRussianAI/OpenAirAI-X"
tokenizer = AutoTokenizer.from_pretrained(model_id)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device).eval()
print(f"✅ Модель загружена на {device}")

# === ХРАНИЛИЩЕ ИСТОРИИ ===
HISTORY_DIR = Path("chat_history")
HISTORY_DIR.mkdir(exist_ok=True)

def get_history_path(username):
    safe_name = "".join(c for c in username if c.isalnum() or c in ('-', '_')).lower()
    return HISTORY_DIR / f"{safe_name}.json"

def load_history(username):
    path = get_history_path(username)
    if path.exists():
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

def save_history(username, history_data):
    path = get_history_path(username)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(history_data, f, ensure_ascii=False, indent=2)

def check_if_pro(token):
    if not token or not token.startswith("hf_"):
        return False
    try:
        response = requests.get(
            "https://huggingface.co/api/whoami-v2",
            headers={"Authorization": f"Bearer {token}"},
            timeout=5
        )
        if response.status_code == 200:
            data = response.json()
            return data.get("isPro", False) or data.get("plan", {}).get("name") == "PRO"
    except Exception as e:
        print(f"Ошибка проверки PRO: {e}")
    return False

# === ОЧИСТКА ОТ ЦИКЛОВ И СПАМА ===
def clean_repetitive_text(text):
    # Убираем явные повторы слов подряд
    words = text.split()
    if len(words) < 4:
        return text
    
    for i in range(len(words) - 3):
        if words[i] == words[i+1] == words[i+2] == words[i+3]:
            return " ".join(words[:i]) + "."
            
    # Если текст содержит подозрительные спам-фразы, можно их обрезать
    spam_triggers = ["mail@gmail.com", "O-pay", "переводить деньги"]
    for trigger in spam_triggers:
        if trigger in text:
            # Обрезаем текст до этого момента
            idx = text.find(trigger)
            return text[:idx].strip() + "."
            
    return text

# === ГЕНЕРАЦИЯ ОТВЕТА ===
def generate_response(message, history, username, current_chat_id):
    if not username:
        gr.Warning("Сначала введите имя пользователя!")
        return history, gr.update(), current_chat_id, gr.update()
    
    if not message.strip():
        return history, gr.update(), current_chat_id, gr.update()
    
    history = history + [{"role": "user", "content": message}]
    
    # ДОБАВЛЯЕМ СИСТЕМНЫЙ ПРОМПТ В НАЧАЛО КАЖДОГО ДИАЛОГА
    # Это заставляет модель помнить, кто она такая
    system_prompt = "Ты — полезный русскоязычный AI-ассистент OpenAirAI. Отвечай кратко, вежливо и по делу. Не упоминай электронную почту, платежи или сторонние организации.\n\n"
    
    prompt = system_prompt
    for msg in history:
        if msg["role"] == "user":
            prompt += f"Пользователь: {msg['content']}\n"
        elif msg["role"] == "assistant":
            prompt += f"AI: {msg['content']}\n"
    prompt += "AI:"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=120,
            min_new_tokens=5,
            temperature=0.9,          # Повышаем температуру для большей вариативности
            top_k=50,                 # Ограничиваем выборку 50 самыми вероятными токенами
            top_p=0.95,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.3    # Штраф за повторы
        )
    
    generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
    response_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
    
    # Очистка
    stop_words = ["Пользователь:", "User:", "\n\n"]
    ai_response = response_text
    for stop_word in stop_words:
        if stop_word in ai_response:
            ai_response = ai_response.split(stop_word)[0].strip()
            
    ai_response = clean_repetitive_text(ai_response)
    
    if not ai_response:
        ai_response = "Извините, я не смог сформулировать ответ."

    history = history + [{"role": "assistant", "content": ai_response}]
    
    if username and current_chat_id:
        history_data = load_history(username)
        if current_chat_id not in history_data:
            history_data[current_chat_id] = {
                "title": message[:40] + ("..." if len(message) > 40 else ""),
                "created": datetime.now().isoformat(),
                "messages": []
            }
        history_data[current_chat_id]["messages"] = history
        save_history(username, history_data)
    
    chat_list_choices = get_chat_list(username)
    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)

# === УПРАВЛЕНИЕ ЧАТАМИ ===
def new_chat(username):
    if not username:
        return [], None, gr.update(choices=[])
    chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    history_data = load_history(username)
    history_data[chat_id] = {
        "title": "Новый чат",
        "created": datetime.now().isoformat(),
        "messages": []
    }
    save_history(username, history_data)
    chat_list_choices = get_chat_list(username)
    return [], chat_id, gr.update(choices=chat_list_choices, value=None)

def load_chat(chat_title, username, current_chat_id):
    if not username or not chat_title:
        return [], current_chat_id
    history_data = load_history(username)
    for cid, data in history_data.items():
        if data["title"] == chat_title:
            return data["messages"], cid
    return [], current_chat_id

def get_chat_list(username):
    if not username:
        return []
    history_data = load_history(username)
    sorted_chats = sorted(
        history_data.items(),
        key=lambda x: x[1].get("created", ""),
        reverse=True
    )
    return [data["title"] for _, data in sorted_chats]

def delete_chat(chat_title, username):
    if not username or not chat_title:
        return [], {}, None, gr.update(choices=[])
    history_data = load_history(username)
    for cid, data in list(history_data.items()):
        if data["title"] == chat_title:
            del history_data[cid]
            break
    save_history(username, history_data)
    chat_list_choices = get_chat_list(username)
    return [], history_data, None, gr.update(choices=chat_list_choices, value=None)

# === CSS ===
CUSTOM_CSS = """
.main-header {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 20px;
    border-radius: 12px;
    margin-bottom: 20px;
    text-align: center;
}
.main-header h1 { margin: 0; font-size: 2em; }
.pro-badge {
    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%);
    color: #333;
    padding: 4px 12px;
    border-radius: 20px;
    font-weight: bold;
    display: inline-block;
    margin-left: 10px;
    box-shadow: 0 2px 8px rgba(255, 215, 0, 0.4);
}
.user-info {
    padding: 10px;
    background: white;
    border-radius: 8px;
    margin-bottom: 10px;
    text-align: center;
    font-weight: 600;
}
.privacy-banner {
    background: linear-gradient(135deg, #10b981 0%, #059669 100%);
    color: white;
    padding: 15px 20px;
    border-radius: 10px;
    margin: 15px 0;
    display: flex;
    align-items: center;
    gap: 12px;
    font-size: 0.95em;
    box-shadow: 0 2px 8px rgba(16, 185, 129, 0.3);
}
.privacy-banner .icon { font-size: 1.8em; }
.privacy-banner strong { display: block; font-size: 1.1em; margin-bottom: 4px; }
"""

PRIVACY_BANNER = """
<div class="privacy-banner">
    <div class="icon">🔒</div>
    <div>
        <strong>Ваши данные в безопасности</strong>
        Все сообщения обрабатываются локально. 
        <b>Данные НЕ передаются в OpenRussianAI</b>. 
        История хранится только внутри контейнера.
    </div>
</div>
"""

# === ИНТЕРФЕЙС ===
with gr.Blocks(title="OpenAirAI-X Chat") as demo:
    username_state = gr.State("")
    history_state = gr.State({})
    current_chat_id_state = gr.State(None)
    
    with gr.Column(visible=False) as main_interface:
        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Русскоязычный AI-ассистент</p></div>")
        gr.HTML(PRIVACY_BANNER)
        
        with gr.Row():
            with gr.Column(scale=1, min_width=250):
                user_info = gr.Textbox(label="Пользователь", value="Не авторизован", interactive=False, elem_classes="user-info")
                pro_badge = gr.HTML('<div class="pro-badge">👑 PRO</div>', visible=False)
                new_chat_btn = gr.Button("➕ Новый чат", variant="primary")
                gr.Markdown("### 📚 История чатов")
                chat_list = gr.Dropdown(choices=[], label="Ваши чаты", interactive=True, allow_custom_value=False, value=None)
                delete_chat_btn = gr.Button("🗑️ Удалить выбранный чат", variant="stop", size="sm")
                logout_btn = gr.Button("🚪 Выйти", size="sm")
            
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(label="Диалог", height=500)
                with gr.Row():
                    msg_input = gr.Textbox(placeholder="Напишите сообщение...", lines=2, scale=5, show_label=False)
                    send_btn = gr.Button("📤 Отправить", variant="primary", scale=1)
    
    with gr.Column(visible=True) as login_screen:
        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Введите данные для входа</p></div>")
        gr.HTML(PRIVACY_BANNER)
        gr.Markdown("### 🔐 Вход в систему\nВведите ваше имя пользователя. История чатов будет привязана к этому имени.")
        username_input = gr.Textbox(label="Имя пользователя (HF Username)", placeholder="Например: RootLinux21")
        token_input = gr.Textbox(label="HF Token (необязательно, для PRO)", placeholder="hf_...")
        login_btn = gr.Button("🔑 Войти", variant="primary")
    
    # === ОБРАБОТЧИКИ ===
    
    def handle_login(username, token):
        if not username.strip():
            gr.Warning("Введите имя пользователя!")
            return gr.update(), gr.update(), "Не авторизован", "", {}, None, gr.update(choices=[], value=None), gr.update(visible=False)
        
        is_pro = check_if_pro(token)
        history_data = load_history(username)
        chat_list_choices = get_chat_list(username)
        
        current_chat_id = None
        if not history_data:
            chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
            history_data[chat_id] = {
                "title": "Новый чат",
                "created": datetime.now().isoformat(),
                "messages": []
            }
            save_history(username, history_data)
            current_chat_id = chat_id
            chat_list_choices = get_chat_list(username)
        
        pro_badge_html = '<div class="pro-badge">👑 PRO</div>'
        display_name = f"👤 {username}"
        
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            display_name,
            username,
            history_data,
            current_chat_id,
            gr.update(choices=chat_list_choices, value=None),
            gr.update(visible=is_pro, value=pro_badge_html)
        )
    
    def handle_logout():
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            "Не авторизован",
            "",
            {},
            None,
            gr.update(choices=[], value=None),
            gr.update(visible=False)
        )
    
    login_btn.click(
        fn=handle_login,
        inputs=[username_input, token_input],
        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
    )
    
    send_btn.click(
        fn=generate_response,
        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
    )
    
    msg_input.submit(
        fn=generate_response,
        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
    )
    
    new_chat_btn.click(
        fn=new_chat,
        inputs=[username_state],
        outputs=[chatbot, current_chat_id_state, chat_list]
    )
    
    chat_list.change(
        fn=load_chat,
        inputs=[chat_list, username_state, current_chat_id_state],
        outputs=[chatbot, current_chat_id_state]
    )
    
    delete_chat_btn.click(
        fn=delete_chat,
        inputs=[chat_list, username_state],
        outputs=[chatbot, history_state, current_chat_id_state, chat_list]
    )
    
    logout_btn.click(
        fn=handle_logout,
        inputs=None,
        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
    )

demo.launch(server_name="0.0.0.0", server_port=7860, css=CUSTOM_CSS)