# nahj engine

import json
import os
import numpy as np
import torch
import faiss
from typing import List, Tuple
# import fasttext
# import fasttext.util
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import datetime
import re
import random
# from embedder import PersianVectorAnalyzer
from normalizer import cleaning
from openai import OpenAI
from langchain_openai import ChatOpenAI
import requests
from query_analyzer import build_messages
import data_model as dm

today = f'{datetime.datetime.now().year}{datetime.datetime.now().month}{datetime.datetime.now().day}'

# -------------------
# مدل‌ها و مسیر داده
# -------------------
# MODEL_PATH = "/home/sabr/MODELS/models--sentence-transformers--paraphrase-multilingual-MiniLM-L12-v2"  
MODEL_PATH = "./models/models--sentence-transformers--paraphrase-multilingual-MiniLM-L12-v2"  
RERANKER_MODEL = "BAAI/bge-reranker-v2-m3"
# RERANKER_MODEL = "/home/gpu/HFHOME/hub/models/BAAI/bge-reranker-v2-m3"
FAISS_INDEX_PATH = "./data-faiss/faiss_index_nahj.index"
FAISS_METADATA_PATH = "./data-faiss/faiss_index_nahj_metadata.json"

RERANK_BATCH = int(os.environ.get("RERANK_BATCH", 256))

def get_key():
    key = 'aa-Fu5oeQv8jx8NCWV39WenJ7Yy1mbcFJ4P20CLQURkql2Eleta' # nahj_rag api key
    return key

def read_records():
    metadata = dm.get_all_data()
    ids, part_ids, context_ids, large_titles, normalized_sentences, titles, urls, arabic_texts, Interpretation_links, types = [], [], [], [], [], [], [], [], [], []
    for item in metadata:
        
        ids.append(item['id'])
        context_ids.append(item['context_id'])
        part_ids.append(item['part_id'])
        titles.append(item['title'])
        large_titles.append(item['large_title'])
        normalized_sentences.append(item['normalized_sentence'])
        urls.append(item['url'])
        types.append(item['types'])
        arabic_texts.append(item['arabic_text'])
        Interpretation_links.append(item['interpretation_links'])
      
        # ids.append(item["id"])
        # part_ids.append(item["part_id"])
        # context_ids.append(item["context_id"])
        # large_titles.append(item["large_title"])
        # normalized_sentences.append(item["normalized_text"])
        # titles.append(item["title"])
        # urls.append(item["url"])
        # arabic_texts.append(item["arabic_text"])
        # Interpretation_links.append(item["Interpretation_link"])
        # types.append(item["type"])
        
    return ids, part_ids,context_ids,large_titles,normalized_sentences, titles, urls, arabic_texts, Interpretation_links, types

def load_faiss_index(index_path: str, metadata_path: str):
    index = faiss.read_index(index_path)

    ids, part_ids,context_ids,large_titles,normalized_sentences, titles, urls, arabic_texts, Interpretation_links, types = read_records()
    
    return ids, part_ids,context_ids,large_titles,normalized_sentences, titles, urls, arabic_texts, Interpretation_links, types, index
        
    

def get_client():
    url = "https://api.avalapis.ir/v1" #avalapis #avalapis
    client = OpenAI(
        api_key=get_key(),
        base_url=url,
    )
    return client

def llm_free_request(llm_messages, model="gemini-2.5-flash-lite"):
    messages.clear()
    print(f'llm_free_request func ...')
    print(f'using model: {model}')
    
    try:
        response = client.chat.completions.create(
            messages=llm_messages,
            model=model,
        )
        answer = response.choices[0].message.content
        
    except Exception as error:
      
        try:
            query = llm_messages[0]['content']
            dm.insert_error(query= query, error_message= error)
            
        except:
            query = '???'
            dm.insert_error(query= query, error_message= error)

        answer = 'متاسفانه خطایی رخ داده است.'
        
    return answer

def llm_request(query, model="gemini-2.5-flash-lite"):
    messages.clear()
    print(f'using model: {model}')
    if query == '':
        return 'لطفا متن سوال را وارد نمائید'
    
    # determine_refrence = """شناسه هر سخنرانی در ابتدای آن و با فرمت "id: {idvalue}" آمده است..."""
    try:
        messages.append({"role": "user", "content": query})
        response = client.chat.completions.create(
            messages= messages,
            model= model,
        )
        answer = response.choices[0].message.content
        # messages.append({"role": "assistant", "content": answer})
    except Exception as error:
        # should write in DATABASE
        dm.insert_error(query= query, error_message= error)

        answer = 'با عرض پوزش؛ متاسفانه خطایی رخ داده است.'
        
    return answer

class QueryAnalysisPipeline:

    def __init__(self):
        """
        llm_client → شیء کلاینت مدل زبانی شما (هر چیزی باشد)
        model_name → برای شمارش توکن‌ها و ارسال
        """
        # self.llm = llm_client
        # self.model_name = model_name

    async def analyze(self, user_query: str):
        """
        مرحله ۱: تحلیل کوئری (سلام + زیربخش‌ها)
        خروجی JSON تحلیل شده را برمی‌گرداند
        """
        messages = build_messages(
        user_query=user_query,
        max_tokens=1024
        )

        result = llm_free_request(messages)

        # پردازش JSON خروجی
        try:
            parsed = json.loads(result)
        except Exception:
            # تلاش دوم → استخراج JSON از متن
            try:
                import re
                json_text = re.search(r"\{.*\}", result, re.S).group(0)
                parsed = json.loads(json_text)
            except Exception as e:
                parsed = {
                    "greeting_reply": "",
                    "sub_questions": [],
                    "final_answer_instruction": ""
                }
                print(f'final exception error: {e}')

        return parsed or []

    async def expand_sub_questions(self, sub_questions):
        """
        مرحله ۲: برای هر زیربخش، پاسخ جدا تولید می‌شود
        """
        answers = []
        for i, question in enumerate(sub_questions):
            print(f'answering questions {i}/{len(sub_questions)}')
            # messages = [
            #     {"role": "user", "content": question}
            # ]
            # resp = llm_free_request(messages)
            resp = bale_chat(question)
            answers.append({"question": question, "answer": resp})
        return answers

    async def final_answer(self, main_question, answer_parts):
        """
        مرحله ۳: تولید پاسخ نهایی با استفاده از:
            - سوال اصلی
            - پاسخ‌های جزئی 
        """
        merged = f"""
سوال اصلی:
{main_question}

پاسخ‌های جزئی:
{answer_parts}

لطفاً پاسخ نهایی را فقط بر اساس «پاسخ‌های جزئی» تولید کن.
پاسخ‌های جزئی، برگرفته از متن نهج البلاغه است.
در هر بخش از پاسخ نهایی تولید شده، عنوان منبع را(در صورت وجود در پاسخ‌های جزئی)، ذکر کن.
"""
        messages = [
            {"role": "user", "content": merged}
        ]
        print(f'generating main question based on sub answers ...')
        return llm_free_request(messages)

# -----------------------------
# Hybrid Retriever with FastText
# -----------------------------
class HybridRetrieverReranker:
    __slots__ = (
        "device", "ids","part_ids","context_ids","large_titles","sentences","titles","urls","arabic_texts","Interpretation_links","types", "N", "embedder", "faiss_index", "vectorizer", "tfidf_matrix", "tokenizer", "reranker", "dense_alpha"
    )
    
    # ids, part_ids,context_ids,large_titles,sentences , titles, urls, arabic_texts, Interpretation_links, types,
    def __init__(self, ids: List[str], part_ids: List[str], context_ids: List[str], large_titles: List[str], sentences: List[str], titles: List[str], urls: List[str], arabic_texts: List[str], Interpretation_links: List[str], types: List[str], faiss_index,
                 dense_alpha: float = 0.6, device: str = None):

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device

        self.ids = ids
        self.part_ids = part_ids
        self.context_ids = context_ids
        self.large_titles = large_titles
        self.sentences = sentences
        self.titles = titles
        self.urls = urls
        self.arabic_texts = arabic_texts
        self.Interpretation_links = Interpretation_links
        self.types = types
        self.faiss_index = faiss_index
        self.N = len(sentences)

        # --- Dense Embedder ---
        print("Loading SentenceTransformer model ...")
        self.embedder = SentenceTransformer(MODEL_PATH, device=self.device)
        # embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device= self.device)
        # embedder.save(MODEL_PATH)
        
        # --- Sparse (TF-IDF) ---
        self.vectorizer = TfidfVectorizer(
            analyzer="word",
            ngram_range=(1, 2),
            token_pattern=r"(?u)\b[\w\u0600-\u06FF]{2,}\b",
        )
        self.tfidf_matrix = self.vectorizer.fit_transform(self.sentences)

        # --- Reranker ---
        # self.tokenizer = AutoTokenizer.from_pretrained(RERANKER_MODEL, use_fast=True, local_files_only= True)
        # self.reranker = AutoModelForSequenceClassification.from_pretrained(
        #     RERANKER_MODEL, local_files_only= True
        # ).to(self.device)

        self.dense_alpha = float(dense_alpha)

    # --- Dense retrieval using FastText ---
    # def _embed_sentence(self, text: str) -> np.ndarray:
    #     tokens = text.split()
    #     vectors = [self.embedder.get_word_vector(tok) for tok in tokens if tok.strip()]
    #     if not vectors:
    #         return np.zeros((self.embedder.get_dimension(),), dtype=np.float32)
    #     return np.mean(vectors, axis=0).astype(np.float32)
    
    def _embed_sentence(self, text: str) -> np.ndarray:
        embedding = self.embedder.encode(
            text,
            convert_to_numpy=True,
            normalize_embeddings=True  # مهم برای cosine similarity
        )
        return embedding.astype(np.float32)

    def dense_retrieve(self, query: str, top_k: int):
        if top_k <= 0:
            return [], np.array([], dtype=np.float32)
        q_emb = self._embed_sentence(query)
        D, I = self.faiss_index.search(np.expand_dims(q_emb, axis=0), top_k)
        return I[0].tolist(), D[0]

    # --- Sparse ---
    def sparse_retrieve(self, query: str, top_k: int):
        if top_k <= 0:
            return [], np.array([], dtype=np.float32)
        k = min(top_k, self.N)
        q_vec = self.vectorizer.transform([query])
        sims = cosine_similarity(q_vec, self.tfidf_matrix).ravel()
        idx = np.argpartition(-sims, kth=k - 1)[:k]
        idx = idx[np.argsort(-sims[idx], kind="mergesort")]
        return idx.tolist(), sims[idx]

    # --- Normalization ---
    @staticmethod
    def _minmax_norm(arr: np.ndarray) -> np.ndarray:
        if arr.size == 0:
            return arr
        a_min = arr.min()
        a_max = arr.max()
        rng = a_max - a_min
        if rng < 1e-12:
            return np.zeros_like(arr)
        return (arr - a_min) / rng

    # --- Fusion (RRF) ---
    def fuse(self, d_idx, d_scores, s_idx, s_scores, top_k=50, k_rrf=60):
        combined = {}
        for rank, idx in enumerate(d_idx):
            score = 1.0 / (k_rrf + rank)
            combined[idx] = combined.get(idx, 0) + score
        for rank, idx in enumerate(s_idx):
            score = 1.0 / (k_rrf + rank)
            combined[idx] = combined.get(idx, 0) + score
        sorted_items = sorted(combined.items(), key=lambda x: x[1], reverse=True)
        cand_idx = [item[0] for item in sorted_items[:top_k]]
        return cand_idx

    # --- Rerank ---
    def rerank(self, query: str, candidate_indices: List[int], passages: List[str], final_k: int) -> List[Tuple[int, float]]:
        if final_k <= 0 or not candidate_indices:
            return []
        texts = [query] * len(candidate_indices)
        pairs = passages
        scores: List[float] = []

        def _iter_batches(max_bs: int):
            bs = max_bs
            while bs >= 16:
                try:
                    with torch.inference_mode():
                        for start in range(0, len(pairs), bs):
                            batch_texts = texts[start:start + bs]
                            batch_pairs = pairs[start:start + bs]
                            inputs = self.tokenizer(
                                batch_texts,
                                batch_pairs,
                                padding=True,
                                truncation=True,
                                max_length=512,
                                return_tensors="pt",
                            ).to(self.device)
                            logits = self.reranker(**inputs).logits.view(-1)
                            scores.extend(logits.detach().cpu().tolist())
                    return True
                except torch.cuda.OutOfMemoryError:
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    bs //= 2
            return False

        success = _iter_batches(max_bs=64)
        if not success:
            raise RuntimeError("Reranker failed due to CUDA OOM.")
        reranked = sorted(
            zip(candidate_indices, scores),
            key=lambda x: x[1],
            reverse=True
        )[:final_k]
        return reranked

    def get_passages(self, cand_idx, sentences_list):
        passages = [sentences_list[idx] for idx in cand_idx]
        return passages

    # --- Search ---
    def search(self, query: str, sentence_list, topk_dense=50, topk_sparse=50,
               pre_rerank_k=50, final_k=10):
        d_idx, d_scores = self.dense_retrieve(query, topk_dense)
        s_idx, s_scores = self.sparse_retrieve(query, topk_sparse)
        # print('---------- dddd scores ----------')
        # for item in d_scores:
        #     print(item)
        # print("---------- ssss scores ----------")
        # for item in s_scores:
        #     print(item)
        pre_rerank_k = final_k
        cand_idx = self.fuse(d_idx, d_scores, s_idx, s_scores, pre_rerank_k)

        # passages = self.get_passages(cand_idx, sentence_list)

        # reranked = self.rerank(query, cand_idx, passages, final_k)
        # return [
        #     {"idx": i, "content": self.sentence_list[i], "rerank_score": score}
        #     for i, score in reranked
        # ]
        
        
        # return [
        #     {"idx": i, "content": self.sentence_list[i], "rerank_score": score}
        #     for i, score in reranked
        # ]
        return [
            {"idx": i, "content": self.sentences[i]}
            for i in cand_idx
        ]

def single_query(query: str):

    query = cleaning(query)
        
    retrived_sections = pipe.search(query, sentences, topk_dense=100, topk_sparse=100, pre_rerank_k=100, final_k=10)
    
    retrived_sections_list = []
    final_similars_text = ''
    id_list = ''
    for i, row in enumerate(retrived_sections):
        # نادیده گرفتن جملات خیلی کوچک
        # if row['content'] and len(row['content'].split()) < 15:
        #     continue
        part_id = part_ids[row['idx']]
        row["part_id"] = part_id
        row["url"] = urls[row['idx']]
        
        title_value = '{' + str(large_titles[row['idx']]) + '}'
        result = f"{i+1}. عنوان: {title_value}\n{row['content']}\n\n"
        final_similars_text += ''.join(result)
        id_list += f"{ids[row['idx']]}\n"
        retrived_sections_list.append(row)
    return id_list, final_similars_text, retrived_sections_list

def find_refrences(llm_answer: str) -> List[str]:
    """
    شناسایی شناسه هایی که مدل زبانی، برای تهیه پاسخ از آنها استفاده کرده است
    
    Args:
        llm_answer(str): متنی که مدل زبانی تولید کرده است
        
    Returns:
        refrence_ids(List[str]): لیستی از شناسه های تشخیص داده شده
    """
    pattern = r"\{[^\}]+\}"
    refrence_ids = re.findall(pattern, llm_answer)
    
    return refrence_ids

def replace_refrences(llm_answer: str, refrences_list:List[str]) -> List[str]:
    """
    شناسایی شناسه هایی که مدل زبانی، برای تهیه پاسخ از آنها استفاده کرده است
    
    Args:
        llm_answer(str): متنی که مدل زبانی تولید کرده است
        refrences_list(List[str]): لیست شناسه ماده های مورد استفاده در پاسخ مدل زبانی
    Returns:
        llm_answer(str), : متن بازسازی شده پاسخ مدل زبانی که شناسه ماده های مورد استفاده در آن، اصلاح شده است
    """
    refrences = ''
    for index, ref in enumerate(refrences_list,1):
            # breakpoint()
            llm_answer = llm_answer.replace(ref, f'[{index}]')
            id = ref.lstrip('{')
            id = id.rstrip('}')
            refrences += ''.join(f'[{index}] https://majles.tavasi.ir/entity/detail/view/qsection/{id}\n')
        
    llm_answer = f'{llm_answer}\n\nمنابع پاسخ‌:\n{refrences.strip()}'
    return llm_answer

# load basic items
ids, part_ids, context_ids, large_titles, sentences, titles, urls, arabic_texts, Interpretation_links, types, faiss_index = load_faiss_index(FAISS_INDEX_PATH, FAISS_METADATA_PATH)

pipe = HybridRetrieverReranker(ids, part_ids,context_ids,large_titles,sentences , titles, urls, arabic_texts, Interpretation_links, types, faiss_index, dense_alpha=0.6)
# query preprocess and normalize
# normalizer_obj = PersianVectorAnalyzer()

messages = [
        {"role": "system", "content": "تو یک دستیار خبره برای کار با محتوای نهج‌البلاغه هستی. پاسخ ها باید الزاما به زبان فارسی باشد. پاسخ ها فقط از متونی که در پرامپت ارائه می شود استخراج شود و از هیچ منبع دیگری استفاده نشود."},
]
client = get_client()
models = [ "gemini-2.5-flash-lite", "gpt-4o-mini","deepseek-reasoner"]

def save_result(chat_obj: object) -> bool:
    # index result in DATABASE
    pass

def run_chatbot(query:str, chat_id:str):
    prompt_status = True
    status_text = 'لطفا متن سوال را وارد نمائید'
    if query == '':
        prompt_status = False
    
    start_time = (datetime.datetime.now())

    # در صورتی که وضعیت پرامپت معتبر باشد، وارد فرایند شو
    if prompt_status:
        idlist, result_passages_text, result_passages_ids = single_query(query)
        end_retrive = datetime.datetime.now()
        print('-'*40)
        retrive_duration = (end_retrive - start_time).total_seconds()
        print(f'retrive duration: {str(retrive_duration)}')
        
        prompt = f'''برای پرسش "{query}" از میان متن های بیانات رهبر معظم انقلاب، پاسخ مناسب را استخراج کن. پاسخ به صورت تحلیلی باشد و ابعاد مختلف پرسش را در نظر بگیرد. پاسخ تولید شده، باید متن های مرتبط با پرسش را به صورت علمی بازنویسی کند. برای هر بخش از متن که در تولید پاسخ استفاده می کنی، عنوان آن بخش و تاریخ را نیز در متن اضافه کن. متن بیانات رهبر معظم انقلاب اسلامی: "{result_passages_text}"'''
        
        llm_model = ''
        for model in models:
            try:
                llm_model = model
                llm_answer = llm_request(prompt, model)
            except Exception as error:
                error = f'model: {model} \n{error}\n\n'
                prompt_status = False
                status_text = 'با عرض پوزش، سرویس موقتا در دسترس نیست. لطفا دقایقی دیگر دوباره تلاش نمائید!'
            
    else:
        chat_obj = {
        'id' : chat_id, # str
        'title' : '', # str
        'user_id' : '',
        'user_query' : query, # str
        'model_key' : llm_model, # str
        'retrived_passage' : result_passages_text, # str
        'retrived_ref_ids' : result_passages_ids, # list[obj]
        'prompt_type' : 'question-answer', # str
        'retrived_duration' : retrive_duration, # str
        'llm_duration' : '0', # str
        'full_duration' : '0', # str
        'time_create' : str(start_time), # str
        'used_ref_ids' : [], # list[str]
        'prompt_answer' : '', # str
        'status' : prompt_status, # or False # bool
        }
        
        # آبجکت ایجاد شده با بازگردان
        return chat_obj, status_text
                
    llm_answer_duration = (datetime.datetime.now() - end_retrive).total_seconds()
    print(f'llm answer duration: {str(llm_answer_duration)}')

    used_refrences_in_answer = find_refrences(llm_answer)
    llm_answer = replace_refrences(llm_answer, used_refrences_in_answer)
            
    full_prompt_duration = (datetime.datetime.now() - start_time).total_seconds()
    print(f'full prompt duration: {full_prompt_duration}')
    print('~'*40)

    chat_obj = {
        'id' : chat_id, # str
        'title' : '', # str
        'user_id' : '',
        'user_query' : query, # str
        'model_key' : llm_model, # str
        'retrived_passage' : result_passages_text, # str
        'retrived_ref_ids' : result_passages_ids, # list[obj]
        'prompt_type' : 'question-answer', # str
        'retrived_duration' : retrive_duration, # str
        'llm_duration' : llm_answer_duration, # str
        'full_duration' : full_prompt_duration, # str
        'time_create' : str(start_time), # str
        'used_ref_ids' : used_refrences_in_answer, # list[str]
        'prompt_answer' : llm_answer, # str
        'status' : True, # or False # bool
    }

    # save_result(chat_obj)
    
    status_text ='پاسخ با موفقیت ایجاد شد'
    return chat_obj, status_text

# @chatbot.post("/credit_refresh")
def credit_refresh():
    """
    بازگرداندن میزان شارژ باقیمانده سامانه
    """
    url = "https://api.avalai.ir/user/credit"
    headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {get_key()}"
    }
    remained_credit = str((requests.get(url, headers=headers)).json()['remaining_irt'])
    
    dm.add_credit(remained_credit= remained_credit)

    return remained_credit
# تعریف مدل داده‌ها برای درخواست‌های API
# class Query(BaseModel):
#     query: str

date = str((datetime.datetime.now())).replace(' ','-').replace(':','').replace('.','-')
chat_id = f'{date}-{random.randint(100000, 999999)}'
print('#'*29)
print('  - NAHJ ENGINE IS READY! - ')
print('#'*29)
# مسیر API برای اجرا کردن run_chatbot
# @chatbot.post("/run_chatbot")
def chat(query):
    
    answer = run_chatbot(query.query, chat_id)
    credit_refresh()

    return {"answer": answer}

pipe_qa = QueryAnalysisPipeline()

async def bale_complex_chat(user_query: str):
    # مرحله 1
    analysis = await pipe_qa.analyze(user_query)
    print("Analysis:", analysis)

    # مرحله 2
    sub_answers = await pipe_qa.expand_sub_questions(analysis["sub_questions"])
    sub_qa = []
    for i, item in enumerate(analysis["sub_questions"]):
        sub_qa.append({
            'question': item,
            'answer': sub_answers[i]['answer']
        })
    print(f'sub answers count: {len(sub_answers)}')
    
    answers_text = ''.join(f'{ans['answer']}\n\n' for ans in sub_answers)
    # مرحله 3
    final_answer = await pipe_qa.final_answer(user_query, answers_text)
    final_result = {
        'final_answer': final_answer,
        'sub_qa': sub_qa
    }
    print(f'len messages COMPLEX: {len(messages)}')
    return final_result

def get_passages_by_paragraphs(retrived_sections_list):
    """
    بازسازی متن های مشابه بر اساس پاراگراف آنها
    """
    data = dm.get_all_data()
    
    final_passages = ''      
    for item in retrived_sections_list:
        filtered_data = {}
        for row in data:
            if row['part_id'] == item['part_id']:
                filtered_data[row['id']] = row
                title = row['large_title']
                
        # مرتب سازی بر اساس ترتیب جمله در پاراگراف
        sorted_data = dict(sorted(filtered_data.items()))

        paragraph = f'{title}:\n'
        for key, value in sorted_data.items():
            paragraph += ''.join(f'{value['normalized_sentence']}. ')
            
        final_passages += ''.join(f'{paragraph.strip()}\n\n')
        
    return final_passages
    

def bale_search(query):
    start = datetime.datetime.now()
    id_list, result_passages, retrived_sections_list  = single_query(query)
    related_paragraphs = get_passages_by_paragraphs(retrived_sections_list)
    
    end_retrive = datetime.datetime.now()
    print('-'*40)
    print(f'retrive duration: {(end_retrive - start).total_seconds()}')
    # پاسخ حداکثر 300 کلمه باشد.
   
    refrences = ''

    print('---------------------------------------------')
    print(f'full duration: {(datetime.datetime.now() - start).total_seconds()}')
    print('---------------------------------------------')
 
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

    return related_paragraphs.strip()

def bale_chat(query):
    start = datetime.datetime.now()
    idlist, result_passages, retrived_sections_list  = single_query(query)
    related_paragraphs = get_passages_by_paragraphs(retrived_sections_list)
    end_retrive = datetime.datetime.now()
    print('-'*40)
    print(f'retrive duration: {(end_retrive - start).total_seconds()}')
    # پاسخ حداکثر 300 کلمه باشد.
    # prompt = f'برای پرسش "{query}" از میان متن های "{result_passages}" .پاسخ مناسب را استخراج کن. پاسخ از زبان رهبر انقلاب اسلامی -حضرت آیت الله خامنه‌ای-  بیان می شود. پاسخ به صورت تحلیلی باشد و ابعاد مختلف پرسش را در نظر بگیرد. پاسخ تولید شده، باید متن های مرتبط با پرسش را به صورت علمی بازنویسی کند. برای هر بخش از متن که در تولید پاسخ استفاده می کنی، عنوان آن بخش و تاریخ را نیز حتما در متن اضافه کن. درصورتی که مطلبی مرتبط با پرسش در متن پیدا نشد، فقط پاسخ بده: "متاسفانه در منابع، پاسخی پیدا نشد!"'
    prompt = f'''برای پرسش "{query}" از میان متن های زیر که از نهج‌البلاغه استخراج شده است، پاسخ مناسب را استخراج کن. متن خروجی «بدون مارک داون» و «بدون استایل» باشد. برای هر کدام از بخش‌ها که در تولید پاسخ استفاده می کنی، عنوان آن بخش را نیز در پاسخ اضافه کن. متن‌های مرتبط از کتاب نهج‌البلاغه "{related_paragraphs.strip()}"'''

    llm_answer = llm_request(prompt)# "deepseek-reasoner"
    print('-'*40)
    print(f'llm duration: {(datetime.datetime.now() - end_retrive).total_seconds()}')

    refrences = ''

    print('---------------------------------------------')
    print(f'full duration: {(datetime.datetime.now() - start).total_seconds()}')
    print('---------------------------------------------')
    print(f'len messages SIMPLE question: {len(messages)}')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

    return llm_answer

# uvicorn src.app:app --reload
if __name__ == "__main__":
    
    # query = 'در قانون حمایت از خانواده و جوانی جمعیت چه خدماتی در نظر گرفته شده است؟'
    while True:
        query = input('enter your qustion:')
        if query == '':
            # should write in DATABASE
            with open('./nahj-answer/chat-leader.txt', 'r', encoding='utf-8') as file:
                query = file.read()
        
        start = (datetime.datetime.now())
        # result = test_dataset()
        idlist, result_passages, retrived_sections_list  = single_query(query)
        end_retrive = datetime.datetime.now()
        print('-'*40)
        print(f'retrive duration: {(end_retrive - start).total_seconds()}')
        
        # prompt = f'برای پرسش "{query}" از میان متن های "{result_passages}" .پاسخ مناسب را استخراج کن. پاسخ از زبان رهبر انقلاب اسلامی -حضرت آیت الله خامنه‌ای-  بیان می شود. پاسخ به صورت تحلیلی باشد و ابعاد مختلف پرسش را در نظر بگیرد. پاسخ تولید شده، باید متن های مرتبط با پرسش را به صورت علمی بازنویسی کند. برای هر بخش از متن که در تولید پاسخ استفاده می کنی، عنوان آن بخش و تاریخ را نیز حتما در متن اضافه کن. درصورتی که مطلبی مرتبط با پرسش در متن پیدا نشد، فقط پاسخ بده: "متاسفانه در منابع، پاسخی پیدا نشد!"'
        prompt = f'''برای پرسش "{query}" از میان بیانات رهبر معظم انقلاب، پاسخ مناسب را استخراج کن. پاسخ به صورت تحلیلی باشد و ابعاد مختلف پرسش را در نظر بگیرد. پاسخ تولید شده، باید محتوای بیانات مرتبط با پرسش را به صورت علمی بازنویسی کند. برای هر بخش از بیانات که در تولید پاسخ استفاده می کنی، عنوان آن و تاریخ را نیز در پاسخ اضافه کن. بیانات رهبر معظم انقلاب اسلامی: "{result_passages}"'''
        
        llm_answer = llm_request(prompt)# "deepseek-reasoner"
        
        print('-'*40)
        print(f'llm duration: {(datetime.datetime.now() - end_retrive).total_seconds()}')

        refrences = ''
        # recognized_refrences = find_refrences(llm_answer)
        # llm_answer = replace_refrences(llm_answer, recognized_refrences)
        
        # should write in DATABASE
        with open('./nahj-answer/result-nahj.txt', mode='a+', encoding='utf-8') as file:
            result_message = f'متن پرامپت: {query.strip()}\n\nپاسخ: {llm_answer} \n----------------------------------------------------------\n'
            file.write(result_message)  
            
        # should write in DATABASE
        with open('./nahj-answer/passages-nahj.txt', mode='a+', encoding='utf-8') as file:
            result_message = f'متن پرامپت: {query.strip()}\nمواد مشابه: {result_passages} \n----------------------------------------------------------\n'
            file.write(result_message)
        
        # should write in DATABASE
        with open('./nahj-answer/chat-nahj.txt', mode='w', encoding='utf-8') as file:
            result_message = f'{query.strip()}\n\nپاسخ:\n {llm_answer} \n----------------------------------------------------------\n'
            file.write(result_message)
            
        print('---------------------------------------------')
        print(f'full duration: {(datetime.datetime.now() - start).total_seconds()}')
        print('---------------------------------------------')
        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
