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

# === ЗАГРУЗКА МОДЕЛИ ===
print("🚀 Загрузка модели...")
model_id = "OpenRussianAI/OpenAirAI-X"
tokenizer = AutoTokenizer.from_pretrained(model_id)
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 generate_response(message, history, username, current_chat_id):
    if not username:
        gr.Warning("Сначала введите имя пользователя на экране входа!")
        return history, gr.update(), current_chat_id
    
    if not message.strip():
        return history, gr.update(), current_chat_id
    
    history = history + [{"role": "user", "content": message}]
    
    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=256,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    ai_response = response.split("AI:")[-1].strip()
    
    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)
    
    return history, gr.update(value=""), current_chat_id

# === УПРАВЛЕНИЕ ЧАТАМИ ===
def new_chat(username):
    if not username:
        return [], None, []
    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)
    return [], chat_id, get_chat_list(username)

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, []
    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)
    return [], history_data, None, get_chat_list(username)

# === 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;
}
"""

# === БАННЕР ПРИВАТНОСТИ (HTML) ===
PRIVACY_BANNER = """
<div class="privacy-banner">
    <div class="icon">🔒</div>
    <div>
        <strong>Ваши данные в безопасности</strong>
        Все сообщения обрабатываются локально в этом Space. 
        <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.Radio(choices=[], label="Ваши чаты", interactive=True)
                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("""
        ### 🔐 Вход в систему
        
        Введите ваше имя пользователя. История чатов будет привязана к этому имени и сохранена **только внутри этого Space**.
        """)
        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")
        gr.Markdown("*Если у вас есть PRO-подписка, вставьте ваш токен (Settings -> Access Tokens), чтобы увидеть значок 👑*")
    
    # === ОБРАБОТЧИКИ ===
    
    def handle_login(username, token):
        if not username.strip():
            gr.Warning("Введите имя пользователя!")
            return gr.update(), gr.update(), "Не авторизован", "", {}, 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,
            chat_list_choices,
            gr.update(visible=is_pro, value=pro_badge_html)
        )
    
    def handle_logout():
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            "Не авторизован",
            "",
            {},
            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]
    ).then(
        fn=lambda u: get_chat_list(u),
        inputs=[username_state],
        outputs=[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]
    ).then(
        fn=lambda u: get_chat_list(u),
        inputs=[username_state],
        outputs=[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]
    )

# Запуск с CSS в launch() (как требует Gradio 6.0)
demo.launch(server_name="0.0.0.0", server_port=7860, css=CUSTOM_CSS)