import gradio as gr
from cryptography.fernet import Fernet
import base64
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import os
import time
import datetime

# Хранилище сообщений в памяти
messages = []
GLOBAL_PASSWORD = "default_password"  # Установите общий пароль для всего чата
SALT = os.urandom(16)  # Генерируем общую соль для всех сообщений
MAX_MESSAGES = 100  # Максимальное количество сохраняемых сообщений

def generate_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    return key

def encrypt_message(text):
    key = generate_key(GLOBAL_PASSWORD, SALT)
    f = Fernet(key)
    return f.encrypt(text.encode())

def decrypt_message(encrypted_data):
    key = generate_key(GLOBAL_PASSWORD, SALT)
    f = Fernet(key)
    return f.decrypt(encrypted_data).decode()

def send_message(username, message_text, password):
    global GLOBAL_PASSWORD
    
    if not message_text or not username:
        return "Введите имя и сообщение", get_chat_history(password)
    
    # Если это первое сообщение, устанавливаем пароль
    if not messages and password:
        GLOBAL_PASSWORD = password
    
    # Проверяем пароль для существующего чата
    if messages and password != GLOBAL_PASSWORD:
        return "❌ Неверный пароль чата", get_chat_history(password)
    
    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
    
    try:
        # Подготовка сообщения для шифрования
        full_message = f"{username} ({timestamp}): {message_text}"
        
        # Шифруем и сохраняем сообщение
        encrypted_data = encrypt_message(full_message)
        messages.append(encrypted_data)
        
        # Ограничиваем количество сообщений
        if len(messages) > MAX_MESSAGES:
            messages.pop(0)
        
        return "", get_chat_history(password)
    except Exception as e:
        return f"Ошибка: {str(e)}", get_chat_history(password)

def get_chat_history(password):
    global GLOBAL_PASSWORD
    
    if not messages:
        return "Чат пуст. Отправьте первое сообщение!"
    
    if password != GLOBAL_PASSWORD:
        return "🔒 Введите правильный пароль чата для просмотра сообщений"
    
    try:
        # Расшифровываем и отображаем все сообщения
        decrypted_messages = []
        for msg in messages:
            decrypted = decrypt_message(msg)
            decrypted_messages.append(decrypted)
        
        return "\n".join(decrypted_messages)
    except Exception as e:
        return f"Ошибка при расшифровке: {str(e)}"

def check_for_updates(chat_history, password):
    time.sleep(1)  # Задержка для предотвращения слишком частого обновления
    return get_chat_history(password)

# Создаем интерфейс Gradio
with gr.Blocks(title="Зашифрованный чат (72 часа)") as demo:
    gr.Markdown("## 🔒 Временный зашифрованный чат")
    gr.Markdown("Чат доступен в течение 72 часов или до закрытия Google Colab")
    
    with gr.Row():
        username = gr.Textbox(label="Ваше имя", placeholder="Введите ваше имя...")
        chat_password = gr.Textbox(label="Пароль чата", type="password", 
                                   placeholder="Пароль для шифрования чата...")
    
    chat_display = gr.Textbox(label="Сообщения чата", lines=15, interactive=False)
    
    with gr.Row():
        message_input = gr.Textbox(label="Сообщение", placeholder="Введите сообщение...", lines=2)
        send_btn = gr.Button("Отправить")
    
    error_box = gr.Textbox(label="Статус", interactive=False, visible=True)
    
    refresh_btn = gr.Button("Обновить чат")
    
    # Обработчики событий
    send_btn.click(
        fn=send_message, 
        inputs=[username, message_input, chat_password], 
        outputs=[error_box, chat_display]
    ).then(lambda: "", outputs=message_input)
    
    refresh_btn.click(
        fn=check_for_updates,
        inputs=[chat_display, chat_password],
        outputs=[chat_display]
    )
    
    # Автоматическое обновление
    chat_display.change(
        fn=lambda: time.sleep(0.5),  # Небольшая пауза
        inputs=None,
        outputs=None
    )
    
    gr.Markdown("""
    ### 📝 Инструкция по использованию
    1. **Первый пользователь** устанавливает пароль чата
    2. **Поделитесь ссылкой** и паролем с участниками чата
    3. **Все участники** должны использовать один и тот же пароль
    4. Нажимайте "Обновить чат" чтобы увидеть новые сообщения
    5. История сообщений хранится только в памяти и будет потеряна при закрытии Colab
    
    ### ⚠️ Важно: 
    - Ссылка будет работать не более 72 часов
    - При закрытии вкладки Google Colab чат будет потерян
    - Чат не сохраняется между сессиями
    """)

demo.launch(share=True)