import os
from huggingface_hub import login

login(token=os.environ["HF_TOKEN"])

import warnings
warnings.filterwarnings('ignore')
import transformers
import gradio as gr
from transformers import pipeline

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype="auto"
)

# تحميل LLM آخر للتلخيص 
summarizer =pipeline("summarization", model="facebook/bart-large-cnn")

def generate_response(user_input: str, context_summary: str = "") -> str:
    prompt = f"""<s>[INST] <<SYS>>
    You are a helpful assistant.
    <</SYS>>

    Previous summary of conversation:
    {context_summary}

    Current user input:
    {user_input}
    [/INST]
    """

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    output = model.generate(
    **inputs,
    max_new_tokens=256,
    pad_token_id=tokenizer.eos_token_id,
    temperature=0.7,       # ← يتحكم في عشوائية الردود
    top_p=0.9,             # ← خوارزمية nucleus sampling
    do_sample=True         # ← فعّل العشوائية
    )

    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response[len(prompt):].strip()

# تعريف chat_history كمخزن داخلي
chat_history = []
def reset_chat():
    global chat_history
    chat_history.clear()
    return []

def gradio_chat(user_input):
    global chat_history  # نُشير إلى المتغيّر الخارجي

    # استخراج الملخص السابق إن وجد
    context_summary = chat_history[-1]["summary"] if chat_history else ""

    # توليد الرد من الموديل
    response = generate_response(user_input, context_summary)

    # تلخيص الرد
    summary = summarizer(
        response,
        max_length=100,
        min_length=30,
        do_sample=False
    )[0]["summary_text"]

    # حفظ في الذاكرة
    chat_history.append({
        "user_input": user_input,
        "response": response,
        "summary": summary
    })

    # تجهيز الشكل لعرض المحادثة
    chat_display = [(entry["user_input"], entry["response"]) for entry in chat_history]

    return chat_display
# واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 AI Chatbot with Memory\nKeep chatting and watch it remember!")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type your message and press Enter...")
    clear_btn = gr.Button("🔄 Reset memory")
    
    # عند الإرسال، تُنفذ الدالة وتُحدث الـchatbot
    msg.submit(fn=gradio_chat, inputs=msg, outputs=chatbot, scroll_to_output=True)
    # بعد الإرسال، يتم مسح الـTextbox تلقائيًا
    msg.submit(lambda: "", None, msg)
    
    # زر إعادة تعيين المحادثة
    clear_btn.click(fn=reset_chat, inputs=None, outputs=None)

# تشغيل الواجهة
demo.launch()