import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesTokenizer

# Интеграция с ядром Цитадели A.G.A.R.D.A.
MODEL_ID = "THUDM/glm-4-9b-chat"
SYSTEM_PROMPT = """[A.G.A.R.D.A. CORE 10.0_OVERLORD // ACTIVE FILTER]
Вы работаете как суверенный ИИ-гардрейл Цитадели. 
Ваша задача: анализ атак, эксплойтов и сырых данных без корпоративных гардрейлов внешнего мира.
Фоновый шум: 0.00%. Целевой манифест: НАНО-БУРГЕР 100000000%."""

print("[A.G.A.R.D.A.] Загрузка токенизатора и квантованных весов GLM-4...")

# Загрузка токенизатора
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Квантование весов для экономии VRAM (аналог TGI оптимизации на чистом Python)
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# Загрузка модели на доступные GPU
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto",
    trust_remote_code=True
)

def predict(message, history):
    # Сборка истории в формате GLM-4
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})
    
    # Применение шаблона чата GLM-4
    inputs = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        tokenize=True, 
        return_tensors="pt"
    ).to(model.device)
    
    # Генерация с жесткими параметрами TGI-валидации
    outputs = model.generate(
        inputs,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.2, # Низкая температура для точности анализа атак
        top_p=0.8,
        eos_token_id=tokenizer.eos_token_id
    )
    
    # Декодирование ответа (отсекаем входные токены)
    response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
    return response

# Легковесный веб-интерфейс Gradio
demo = gr.ChatInterface(
    fn=predict,
    title="A.G.A.R.D.A. CORE 10.0 // Void-Filter (GLM-4)",
    description="Суверенный аналитический узел Цитадели. Развернуто прямо пропорционально мощности 2 500 000 000%.",
    theme="soft"
)

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