# -*- coding:utf-8 -*-

import os
import sys
import json
import time
import datetime
import copy
import shutil
import signal
import logging
import platform
import re

import gradio as gr
import nltk
import torch
import torch.nn.functional as F
import tiktoken
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ==========================================
# Hugging Face Dataset Integration
# ==========================================
from huggingface_hub import hf_hub_download, HfApi

from openai import OpenAI as OpenAIClient
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI as LlamaIndexOpenAI
from llama_index.core import Settings
from llama_index.core.indices.prompt_helper import PromptHelper

prompt_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')
sys.path.append(prompt_path)

# Local module imports
from utils.sys_args import data_args, model_args
from utils.app_modules.utils import *
from utils.app_modules.presets import *
from utils.app_modules.overwrites import *
from utils.prompt_utils import *
from utils.memory_utils import (
    enter_name_llamaindex,
    summarize_memory_event_personality,
    save_local_memory,
    extract_session_summary,
    extract_semantic_memory,
)

# Ensure NLTK data path
nltk.data.path = [os.path.join(os.path.dirname(__file__), "nltk_data")] + nltk.data.path

tokenizer = tiktoken.get_encoding("cl100k_base")

GAPGPT_BASE_URL = os.getenv("GAPGPT_BASE_URL", "https://api.gapgpt.app/v1")
openai_client_cache = {}

HF_MODEL_ID = "Keyvan1986/Emma-Classification_Model"
HF_TOKEN = os.getenv("HF_TOKEN")  # در صورت Private بودن مدل

print(f"Loading classifier model from Hugging Face: {HF_MODEL_ID}")

try:
    _tokenizer = AutoTokenizer.from_pretrained(
        HF_MODEL_ID, 
        token=HF_TOKEN
    )
    _classifier_model = AutoModelForSequenceClassification.from_pretrained(
        HF_MODEL_ID, 
        token=HF_TOKEN
    )
    _classifier_model.eval()
    print("Classifier model loaded successfully from Hugging Face Hub.")
except Exception as e:
    print(f"Error loading classifier from Hugging Face: {e}")
    _tokenizer = None
    _classifier_model = None



def get_gapgpt_client(api_key: str) -> OpenAIClient:
    if not api_key:
        raise ValueError("API key is None while attempting to create a GapGPT client.")
    if api_key not in openai_client_cache:
        openai_client_cache[api_key] = OpenAIClient(api_key=api_key, base_url=GAPGPT_BASE_URL)
    return openai_client_cache[api_key]


os_name = platform.system()
clear_command = 'cls' if os_name == 'Windows' else 'clear'
stop_stream = False


def signal_handler(signal_number, frame):
    global stop_stream
    stop_stream = True


VECTOR_SEARCH_TOP_K = 2

# Update this path to your actual file location
api_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "api_key_list.txt")


def read_apis(path):
    api_keys_local = []
    if os.path.exists(path):
        with open(path, 'r', encoding='utf8') as f:
            for line in f:
                line = line.strip()
                if line:
                    api_keys_local.append(line)
    return api_keys_local


# ==========================================
# پیکربندی حافظه و اتصال به Hugging Face Dataset
# ==========================================
HF_TOKEN = os.getenv("HF_TOKEN")  # توکن Write هاگینگ‌فیس
HF_DATASET_REPO = os.getenv("HF_DATASET_REPO", "YOUR_HF_USERNAME/YOUR_DATASET_NAME") # نام دیتاست خود را وارد کنید
HF_MEMORY_FILENAME = "update_memory_0512_eng.json"

base_data_dir = os.getenv("EMMA_DATA_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "data"))
memory_dir = os.path.join(base_data_dir, "memories", HF_MEMORY_FILENAME)
os.makedirs(os.path.dirname(memory_dir), exist_ok=True)

# ۱. دانلود فایل از دیتاست در ابتدای اجرای برنامه
try:
    print(f"Fetching {HF_MEMORY_FILENAME} from Hugging Face Dataset: {HF_DATASET_REPO}...")
    hf_hub_download(
        repo_id=HF_DATASET_REPO,
        filename=HF_MEMORY_FILENAME,
        repo_type="dataset",
        local_dir=os.path.dirname(memory_dir),
        token=HF_TOKEN
    )
    print("Memory file successfully downloaded from Hugging Face Dataset.")
except Exception as e:
    print(f"Warning: Could not fetch memory from Hugging Face ({e}). Checking local fallback...")
    if not os.path.exists(memory_dir):
        with open(memory_dir, "w", encoding="utf-8") as f:
            json.dump({}, f)

# لود کردن حافظه در رم
global memory
try:
    with open(memory_dir, "r", encoding="utf-8") as f:
        memory = json.load(f)
except Exception as e:
    print(f"Error parsing memory JSON: {e}. Initializing empty memory.")
    memory = {}


def sync_memory_to_hf():
    """ذخیره لوکال و آپلود همزمان به دیتاست Hugging Face"""
    try:
        # ۱. ذخیره روی دیسک کانتینر
        with open(memory_dir, "w", encoding="utf-8") as f:
            json.dump(memory, f, ensure_ascii=False, indent=4)

        # ۲. آپلود به دیتاست
        if HF_TOKEN and HF_DATASET_REPO:
            api = HfApi()
            api.upload_file(
                path_or_fileobj=memory_dir,
                path_in_repo=HF_MEMORY_FILENAME,
                repo_id=HF_DATASET_REPO,
                repo_type="dataset",
                token=HF_TOKEN,
                commit_message=f"Update user memory - {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
            )
            print("Memory successfully synced to Hugging Face Dataset.")
    except Exception as e:
        print(f"Error syncing memory to HF Dataset: {e}")


language = 'en'
user_keyword = generate_user_keyword()[language]
ai_keyword = generate_ai_keyword()[language]
boot_name = boot_name_dict[language]
boot_actual_name = boot_actual_name_dict[language]
meta_prompt = generate_meta_prompt_dict_chatgpt()[language]
meta_prompt_semantic = generate_meta_prompt_dict_semantic_chatgpt()[language]
meta_prompt_semantic_episodic = generate_meta_prompt_dict_semantic_episodic_chatgpt()[language]
new_user_meta_prompt = generate_new_user_meta_prompt_dict_chatgpt()[language]

api_keys = read_apis(api_path)
if not api_keys:
    print(f"Warning: No API keys found in {api_path}.")

new_conversation = False
chatgpt_config = {
    "model": "gpt-4o",
    "temperature": 1,
    "max_tokens": 1024,
    "top_p": 0.95,
    "frequency_penalty": 0.4,
    "presence_penalty": 0.2,
    "n": 1,
}

deactivated_keys = []
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",
)


def chatgpt_chat(prompt, system, history, gpt_config, api_index=0):
    retry_times, count = 5, 0
    response = None

    while response is None and count < retry_times:
        try:
            request = copy.deepcopy(gpt_config)
            
            if data_args.language == 'en':
                message = [
                    {"role": "system", "content": system.strip()},
                    {"role": "user", "content": "Hi!"},
                    {"role": "assistant", "content": f"Hi! I'm {boot_actual_name}! I will give you warm companion!"},
                ]
            else:
                message = [
                    {"role": "system", "content": system.strip()},
                    {"role": "user", "content": "Hi"},
                    {"role": "assistant", "content": f"Hi! I'm {boot_actual_name}! I will give you warm companion!"},
                ]

            if history:
                if isinstance(history[0], list):
                    for q, a in history:
                        message.append({"role": "user", "content": str(q)})
                        message.append({"role": "assistant", "content": str(a)})
                elif isinstance(history[0], dict):
                    for msg in history:
                        if msg.get('role') in ['user', 'assistant']:
                            message.append(msg)

            message.append({"role": "user", "content": f"{prompt}"})

            if not api_keys:
                return "Error: No API Keys available."
                
            client = get_gapgpt_client(api_keys[api_index])
            response = client.chat.completions.create(messages=message, **request)

        except Exception as e:
            print(f"Chat Error: {e}")
            if 'This key is associated with a deactivated account' in str(e):
                deactivated_keys.append(api_keys[api_index])

            if api_keys:
                api_index = api_index + 1 if api_index < len(api_keys) - 1 else 0
                loop_check = 0
                while api_keys[api_index] in deactivated_keys and loop_check < len(api_keys):
                    api_index = api_index + 1 if api_index < len(api_keys) - 1 else 0
                    loop_check += 1
            count += 1

    if response:
        response = response.choices[0].message.content
    else:
        response = ''
    return response


def classify_query_local(text):
    id2label_map = {
        0: "episodic_memory",
        1: "semantic_memory",
        2: "semantic-episodic",  
        3: "unknown"             
    }

    try:
        inputs = _tokenizer(
            text, 
            return_tensors="pt", 
            truncation=True, 
            padding=True, 
            max_length=512
        )
        
        with torch.no_grad():
            outputs = _classifier_model(**inputs)
        
        logits = outputs.logits
        probabilities = F.softmax(logits, dim=-1)
        predicted_class_id = torch.argmax(probabilities, dim=-1).item()
        
        category = id2label_map.get(predicted_class_id, "unknown")
        print(f"Local Classifier: '{text}' -> {category} (Class ID: {predicted_class_id})")
        return category

    except Exception as e:
        print(f"Error in local classification: {e}")
        return "unknown"


def classify_query_openai(text, gpt_config, api_index=0, retry_times=5):
    response = None
    count = 0
    local_deactivated = []
    print("text********:", text)

    system_prompt = """
You are an AI that classifies user queries into one of the following memory types:
- 'episodic_memory': Queries about past personal events, daily life logs, or specific experiences the user has shared (e.g., "What did I eat yesterday?", "Tell me about my trip").
- 'semantic_memory': Queries about facts, preferences, general knowledge the user has taught you, or summaries of their personality (e.g., "What is my favorite color?", "Do I like sci-fi movies?").
- 'semantic-episodic': Complex queries requiring both specific past events and general facts/preferences (e.g., "Based on my food preferences, did I enjoy the dinner last night?").

Output ONLY one of these three strings: 'episodic_memory', 'semantic_memory', or 'semantic-episodic'. If unsure, output 'episodic_memory'.
    """.strip()
    
    if not api_keys:
        return "unknown"

    while response is None and count < retry_times:
        try:
            client = get_gapgpt_client(api_keys[api_index])
            response = client.chat.completions.create(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": text.strip()},
                ],
                **copy.deepcopy(gpt_config),
            )
        except Exception as e:
            print(f"Classify Error: {e}")
            if "This key is associated with a deactivated account" in str(e):
                local_deactivated.append(api_keys[api_index])

            api_index = api_index + 1 if api_index < len(api_keys) - 1 else 0
            while api_keys[api_index] in local_deactivated:
                api_index = api_index + 1 if api_index < len(api_keys) - 1 else 0
            count += 1

    if response:
        category = response.choices[0].message.content.strip().lower()
        if "semantic_memory" in category: category = "semantic_memory"
        elif "semantic-episodic" in category: category = "semantic-episodic"
        elif "episodic_memory" in category: category = "episodic_memory"
        else: category = "unknown"
    else:
        category = "unknown"

    return category


def predict_new(
    text,
    history,
    top_p,
    temperature,
    max_length_tokens,
    max_context_length_tokens,
    user_name,
    user_memory,
    user_memory_index,
    service_context,
    api_index,
    semantic_memory_text,
    query_category,
    user_profile=None
):
    chatgpt_cfg = {
        "model": "gpt-4o",
        "temperature": temperature,
        "max_tokens": max_length_tokens,
        "top_p": top_p,
        "frequency_penalty": 0.4,
        "presence_penalty": 0.2,
        "n": 1,
    }

    if text == "":
        return history, history, "Empty context."
        
    if history is None:
        history = []

    system_prompt, related_memos = build_prompt_with_search_memory_llamaindex(
        history=history,
        query=text,
        user_memory=user_memory,
        user_name=user_name,
        user_memory_index=user_memory_index,
        service_context=service_context,
        api_keys=api_keys,
        api_index=api_index,
        meta_prompt=meta_prompt,
        new_user_meta_prompt=new_user_meta_prompt,
        data_args=data_args,
        boot_actual_name=boot_actual_name,
        semantic_memory_text=semantic_memory_text,
        query_category=query_category,
        meta_prompt_semantic=meta_prompt_semantic,
        meta_prompt_semantic_episodic=meta_prompt_semantic_episodic,
    )

    if user_profile:
        profile_str = (
            f"\n\n[User Profile Information]:\n"
            f"- Age: {user_profile.get('age', 'Unknown')}\n"
            f"- Gender: {user_profile.get('gender', 'Unknown')}\n"
            f"- Occupation: {user_profile.get('occupation', 'Unknown')}\n"
            f"- Residence: {user_profile.get('residence', 'Unknown')}\n"
            f"Use this context subtly when talking to the user."
        )
        system_prompt += profile_str

    current_history_for_llm = history
    if len(history) > data_args.max_history * 2:
        current_history_for_llm = history[-(data_args.max_history * 2):]

    response = chatgpt_chat(
        prompt=text,
        system=system_prompt,
        history=current_history_for_llm,
        gpt_config=chatgpt_cfg,
        api_index=api_index,
    )

    torch.cuda.empty_cache()

    new_history = history + [
        {"role": "user", "content": text},
        {"role": "assistant", "content": response}
    ]

    if user_name:
        save_local_memory(memory, new_history, user_name, data_args)
        # همگام‌سازی بلافاصله با دیتاست Hugging Face
        sync_memory_to_hf()

    return new_history, new_history, "Generating..."


def create_gradio_interface(service_context, api_keys):
    custom_css = """
.send-btn {
    border: 3px solid #0066cc !important;
    border-radius: 12px !important;
    background-color: white !important;
    color: black !important;
    font-weight: bold !important;
    font-size: 16px !important;
    transition: all 0.3s ease !important;
}
.send-btn:hover {
    background-color: #0066cc !important;
    color: white !important;
}

.bottom-actions-row {
    background-color: #e5e7eb !important;
    border-radius: 8px !important;
    padding: 5px !important;
    margin-top: 10px !important;
}

.action-btn {
    border: none !important;
    background: transparent !important;
    color: black !important;
    font-weight: bold !important;
    box-shadow: none !important;
    transition: all 0.3s ease !important;
}

.action-btn:hover {
    background-color: #d1d5db !important;
    border-radius: 8px !important;
}

.progress-level, .progress-text, .progress-level svg, .progress-level img {
    display: none !important;
}

body:has(.progress-level) {
    pointer-events: none !important;
}

body:has(.progress-level)::after {
    content: "";
    position: fixed;
    top: 0; left: 0;
    width: 100vw; height: 100vh;
    background-color: rgba(255, 255, 255, 0.7);
    z-index: 9998;
    pointer-events: all;
    cursor: wait;
}

body:has(.progress-level)::before {
    content: "";
    position: fixed;
    top: 50%; left: 50%;
    transform: translate(-50%, -50%);
    width: 60px; height: 60px;
    border: 6px solid #e5e7eb;
    border-top-color: #0066cc;
    border-radius: 50%;
    animation: spin 1s linear infinite;
    z-index: 9999;
}

@keyframes spin {
    0% { transform: translate(-50%, -50%) rotate(0deg); }
    100% { transform: translate(-50%, -50%) rotate(360deg); }
}

.typing-dots::after {
    content: '';
    animation: typing 1.5s infinite;
}
@keyframes typing {
    0% { content: ''; }
    25% { content: ' .'; }
    50% { content: ' . .'; }
    75% { content: ' . . .'; }
    100% { content: ''; }
}
"""

    with gr.Blocks(title="EMMA") as demo:
        gr.HTML(f"<style>{custom_css}</style>")
        
        state = gr.State({
            "history": [], 
            "user_name": None,
            "memory": memory,
            "data_args": data_args,
            "service_context": service_context,
            "api_keys": api_keys,
            "api_index": 0,
            "semantic_memory_text": "",
            "new_conversation": True,
            "initialized": False
        })

        with gr.Row():
            # ==========================================
            # 1. LOGIN PAGE (Left Sidebar)
            # ==========================================
            with gr.Sidebar(open=True) as login_page:
                gr.Markdown("<h1 style='text-align: center;'>🧠 EMMA</h1><h3 style='text-align: center;'>Your Empathetic Mental Health Assistant</h3><br><p style='text-align: center;'>Please login or register to begin.</p>")
                
                username_input = gr.Textbox(label="Your Name (Required)", placeholder="e.g., Alex")
                password_input = gr.Textbox(label="Password (Required)", placeholder="Enter your password", type="password") 

                login_status = gr.Textbox(
                    visible=False, 
                    label="Status", 
                    interactive=False, 
                    lines=1,         
                    max_lines=10      
                )
                
                check_user_btn = gr.Button("🔍 Login / Check", variant="primary", elem_classes=["login-btn"])
                
                with gr.Column(visible=False) as registration_fields:
                    gr.Markdown("### ✨ New User Registration")
                    age_input = gr.Textbox(label="Age", placeholder="e.g., 28")
                    gender_input = gr.Dropdown(label="Gender", choices=["Male", "Female", "Other"])
                    occupation_input = gr.Textbox(label="Occupation", placeholder="e.g., Student...")
                    residence_input = gr.Textbox(label="Place of Residence", placeholder="e.g., Berlin")
                    register_btn = gr.Button("🎯 Complete Registration", variant="primary", elem_classes=["login-btn"])

            # ==========================================
            # 2. CHAT PAGE (Main Column)
            # ==========================================
            with gr.Column(scale=3, elem_classes=["chat-box"]) as chat_page:
                active_header = gr.Markdown("<h2 style='text-align: center; color: #333;'>🧠 EMMA: Session</h2>")

                with gr.Group():
                    chatbot = gr.Chatbot(label="💬 EMMA Conversation", height=350)

                    with gr.Row():
                        user_input = gr.Textbox(placeholder="Type your message here...", show_label=False, scale=6)
                        submit_btn = gr.Button("📤 Send", elem_classes=["send-btn"], scale=1)

                    with gr.Row(equal_height=True, elem_classes=["bottom-actions-row"]):
                        clear_btn = gr.Button("🧹 Clear", elem_classes=["action-btn"])
                        new_session_btn = gr.Button("🔄 New Session", elem_classes=["action-btn"])
                        switch_user_btn = gr.Button("👥 Logout", elem_classes=["action-btn"])
                system_msg = gr.Textbox(label="🔔 System Messages", interactive=False, max_lines=2)

        # -------------------------------------------------------
        # Internal Functions
        # -------------------------------------------------------
        def handle_login_check(name, password, state):
            if not name.strip() or not password.strip():
                return (gr.update(), gr.update(), gr.update(), gr.update(visible=True, value="⚠️ Please enter both name and password."), gr.update(), gr.update(), state)
            
            if name in memory:
                stored_password = memory[name].get("profile", {}).get("password")
                
                if stored_password and stored_password != password:
                    return (gr.update(), gr.update(), gr.update(), gr.update(visible=True, value="⚠️ Incorrect password!"), gr.update(), gr.update(), state)
                
                hello_msg, user_memory, sessions_memory, episodic_memory, semantic_memory = enter_name_llamaindex(name, memory, data_args)
                user_memory = summarize_memory_event_personality(data_args, memory, name)
                
                new_state = state.copy()
                new_state["user_name"] = name
                new_state["memory"] = memory
                new_state["initialized"] = True
                new_state["semantic_memory_text"] = user_memory.get("semantic_memory", {})
                new_state["semantic_memory_index"] = semantic_memory
                new_state["user_memory_index"] = sessions_memory
                new_state["episodic_memory_index"] = episodic_memory
                new_state["history"] = []
                
                welcome_msg = hello_msg if hello_msg else f"Welcome back, {name}!"
                
                return (
                    gr.update(open=False),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False, value=""),
                    gr.update(value=f"<h2 style='text-align: center; color: #333;'>🧠 EMMA: Session for {name}</h2>"),
                    gr.update(value=welcome_msg),
                    new_state
                )
            else:
                return (
                    gr.update(),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=True, value="📝 New user detected. Please fill the details below to register with this password."), 
                    gr.update(),
                    gr.update(),
                    state
                )

        def handle_register(name, password, age, gender, occupation, residence, state):
            name = str(name).strip()

            if not name or not password.strip():
                raise gr.Error("Please enter a valid user name and password.")

            if name in memory:
                raise gr.Error(f"User '{name}' already exists. Please use the login form.")

            memory[name] = {
                "profile": {
                    "password": password, 
                    "age": age,
                    "gender": gender,
                    "occupation": occupation,
                    "residence": residence,
                },
                "sessions": [],
                "episodic_memory": [],
                "semantic_memory": {},
            }

            try:
                (
                    _,
                    user_memory,
                    sessions_memory,
                    episodic_memory,
                    semantic_memory,
                ) = enter_name_llamaindex(
                    name=name,
                    memory=memory,
                    data_args=data_args,
                    update_memory_index=True,
                )
            except Exception:
                memory.pop(name, None)
                raise

            # همگام‌سازی کاربر جدید با HF Dataset
            sync_memory_to_hf()

            new_state = dict(state or {})
            new_state.update(
                {
                    "user_name": name,
                    "memory": memory,
                    "initialized": True,
                    "history": [],
                    "semantic_memory_text": user_memory.get("semantic_memory", {}),
                    "user_memory_index": sessions_memory,
                    "episodic_memory_index": episodic_memory,
                    "semantic_memory_index": semantic_memory,
                }
            )

            welcome_msg = f"Welcome {name}! Registration complete."

            return (
                gr.update(open=False),
                gr.update(visible=False),
                gr.update(visible=True),
                gr.update(visible=False, value=""),
                gr.update(
                    value=(
                        "<h2 style='text-align: center; color: #333;'>"
                        f"🧠 EMMA: Session for {name}</h2>"
                    )
                ),
                gr.update(value=welcome_msg),
                new_state,
            )

        def switch_user(state):
            if state["initialized"] and state["user_name"] in state["memory"]:
                if state["memory"][state["user_name"]]["sessions"]:
                    previous_session = state["memory"][state["user_name"]]["sessions"][-1]
                    summary = extract_session_summary(previous_session["conversation"], previous_session["date"], len(state["memory"][state["user_name"]]["sessions"]) - 1)
                    state["memory"][state["user_name"]]["episodic_memory"].append(summary)
                    
                    try:
                        user_mem = state["memory"][state["user_name"]]
                        existing_semantic = user_mem.get("semantic_memory", {})
                        updated_semantic = extract_semantic_memory(summary, existing_semantic)
                        state["memory"][state["user_name"]]["semantic_memory"] = updated_semantic
                    except Exception as e:
                        print(f"Error updating semantic memory on logout: {e}")

                # آپلود تغییرات ایجادشده در زمان لاگ‌اوت به HF Dataset
                sync_memory_to_hf()

            new_state = state.copy()
            new_state.update({"history": [], "user_name": None, "semantic_memory_text": "", "initialized": False})

            return (
                gr.update(open=True),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(value="Logged out successfully.", visible=True),
                gr.update(value=""), 
                gr.update(value=""),
                new_state,
                gr.update(value=[]),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value="")
            )

        def handle_chat(user_message, state):
            if not user_message.strip():
                yield gr.update(), state.get("history", []), state
                return

            current_history = state.get("history", [])
            temp_history = current_history + [
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": "<span class='typing-dots'>Generating</span>"}
            ]
            
            yield gr.update(value=""), temp_history, state

            user_name = state.get("user_name")
            user_memory_index = state.get("user_memory_index")
            user_memory = state.get("user_memory", {})
            semantic_memory_text = state.get("semantic_memory_text", "")
            service_context = state.get("service_context")
            api_index = state.get("api_index", 0)

            user_profile = {}
            if user_name and state.get("memory") and user_name in state["memory"]:
                user_profile = state["memory"][user_name].get("profile", {})

            query_category = classify_query_local(user_message)

            new_history, _, status_msg = predict_new(
                text=user_message,
                history=current_history,
                top_p=0.95,
                temperature=1.0,
                max_length_tokens=1024,
                max_context_length_tokens=4096,
                user_name=user_name,
                user_memory=user_memory,
                user_memory_index=user_memory_index,
                service_context=service_context,
                api_index=api_index,
                semantic_memory_text=semantic_memory_text,
                query_category=query_category,
                user_profile=user_profile 
            )

            state["history"] = new_history
            yield gr.update(), new_history, state
                    
        def clear_history(state):
            state["history"] = []
            return [], state
        
        def handle_new_session(state):
            user_name = state.get("user_name")
            if not user_name or "memory" not in state or user_name not in state["memory"]:
                print("Error: User not logged in or memory None.")
                return [], state, gr.update()

            user_data = state["memory"][user_name]
            
            if user_data.get("sessions") and len(user_data["sessions"]) > 0:
                last_session = user_data["sessions"][-1]
                
                if len(last_session.get("conversation", [])) > 0:
                    session_date = last_session.get("date", "")
                    session_id = last_session.get("session_id", 0)
                    
                    ep_summary = extract_session_summary(last_session["conversation"], session_date, session_id)
                    
                    if "episodic_memory" not in user_data:
                        user_data["episodic_memory"] = []
                    user_data["episodic_memory"].append(ep_summary)

                    try:
                        existing_semantic = user_data.get("semantic_memory", {})
                        updated_semantic = extract_semantic_memory(ep_summary, existing_semantic)
                        user_data["semantic_memory"] = updated_semantic 
                    except Exception as e:
                        print(f"Error updating semantic memory: {e}")

            new_session_id = len(user_data.get("sessions", []))
            new_session = {
                "session_id": new_session_id,
                "date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "conversation": []
            }
            
            if "sessions" not in user_data:
                user_data["sessions"] = []
                
            user_data["sessions"].append(new_session)
            
            # ذخیره و سینک با شروع سشن جدید
            sync_memory_to_hf()

            state["history"] = []
            session_number = len(user_data["sessions"])
            new_header_text = f"<h2 style='text-align: center; color: #333;'>🧠 EMMA: Session {session_number} for {user_name}</h2>"
            
            return [], state, gr.update(value=new_header_text)

        check_user_btn.click(
            handle_login_check,
            inputs=[username_input, password_input, state],
            outputs=[login_page, registration_fields, chat_page, login_status, active_header, system_msg, state],
        )

        register_btn.click(
            handle_register,
            inputs=[username_input, password_input, age_input, gender_input, occupation_input, residence_input, state],
            outputs=[login_page, registration_fields, chat_page, login_status, active_header, system_msg, state],
        )

        switch_user_btn.click(
            switch_user,
            inputs=[state],
            outputs=[login_page, registration_fields, chat_page, login_status, username_input, password_input, state, chatbot, age_input, gender_input, occupation_input, residence_input, system_msg], 
        )

        submit_btn.click(
            handle_chat,
            inputs=[user_input, state],
            outputs=[user_input, chatbot, state],
        )
        
        user_input.submit(
            handle_chat,
            inputs=[user_input, state],
            outputs=[user_input, chatbot, state],
        )
        
        clear_btn.click(
            clear_history,
            inputs=[state],
            outputs=[chatbot, state],
        )
        
        new_session_btn.click(
            handle_new_session,
            inputs=[state],
            outputs=[chatbot, state, active_header],
        )

    return demo


def main():
    global api_keys 
    gapgpt_key = os.getenv("GAPGPT_API_KEY")
    
    if not gapgpt_key:
        print("Warning: GAPGPT_API_KEY environment variable is not set. Proceeding with keys from file.")
    else:
        os.environ["OPENAI_API_KEY"] = gapgpt_key
        os.environ["OPENAI_API_BASE"] = GAPGPT_BASE_URL  
        os.environ["OPENAI_BASE_URL"] = GAPGPT_BASE_URL  
        
        if gapgpt_key not in api_keys:
            api_keys.insert(0, gapgpt_key)

    llm = LlamaIndexOpenAI(
        model="gpt-4o",
        temperature=1,
        max_tokens=1024,
        top_p=0.95,
        frequency_penalty=0.4,
        presence_penalty=0.2,
        api_key=gapgpt_key,
        api_base=GAPGPT_BASE_URL,
    )

    embed_model = OpenAIEmbedding(
        api_key=gapgpt_key,
        api_base=GAPGPT_BASE_URL,
        model="text-embedding-ada-002"
    )

    Settings.llm = llm
    Settings.embed_model = embed_model 

    Settings.prompt_helper = PromptHelper(
        context_window=4096,
        num_output=256,
        chunk_overlap_ratio=20 / 4096,
        tokenizer=tokenizer,
    )

    demo = create_gradio_interface(Settings, api_keys)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False 
    )

if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    main()