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):
    return HISTORY_DIR / f"{username}.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:
        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("Пожалуйста, войдите через Hugging Face")
        return history, gr.update(), current_chat_id
    
    if not message.strip():
        return history, gr.update(), current_chat_id
    
    # Новый формат Gradio 6+: список словарей
    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;
}
"""

# === ИНТЕРФЕЙС ===
with gr.Blocks(css=CUSTOM_CSS, 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>
        """)
        
        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,
                    type="messages"
                )
                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>Войдите через Hugging Face, чтобы начать</p>
        </div>
        """)
        gr.Markdown("""
        ### 🔐 Вход в систему
        
        Для использования чата необходимо войти через аккаунт Hugging Face.
        
        **Преимущества входа:**
        - 💾 Сохранение истории чатов
        - 👑 PRO-статус для подписчиков
        - 🔄 Доступ к предыдущим диалогам
        """)
        gr.LoginButton(value="🔑 Войти через Hugging Face", variant="primary")
    
    # === ОБРАБОТЧИКИ ===
    
    def handle_login(request: gr.Request):
        if request.username:
            username = request.username
            token = request.oauth_token.get("access_token") if request.oauth_token else None
            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)
            )
        return (gr.update(), gr.update(), gr.update(), "", {}, None, [], gr.update())
    
    def handle_logout():
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            "Не авторизован",
            "",
            {},
            None,
            [],
            gr.update(visible=False)
        )
    
    demo.load(
        fn=handle_login,
        inputs=None,
        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]
    )

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