import sqlite3
import json

# ایجاد جدول (فقط یک بار اجرا می‌شود)
def create_speechs_table():
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    # ایجاد جدول اگر وجود نداشته باشد
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS speeches (
        id TEXT PRIMARY KEY,
        context_id TEXT,
        part_id TEXT,
        title TEXT,
        large_title TEXT,
        normalized_sentence TEXT,
        url TEXT,
        types TEXT,
        arabic_text TEXT,
        interpretation_links TEXT
    )
    ''')
    conn.commit()
    conn.close()

# درج یک لیست از داده‌ها در جدول
def insert_data(data_list):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    for data in data_list:
        try:
            cursor.execute('''
            INSERT INTO speeches (id, context_id, part_id, title, large_title, normalized_sentence, url, types, arabic_text, interpretation_links)
            VALUES (:id, :context_id, :part_id, :title, :large_title, :normalized_sentence, :url, :types, :arabic_text, :Interpretation_links)
            ''', data)
        except sqlite3.IntegrityError:
            print(f"Warning: Data with id '{data['id']}' already exists and was skipped.")

    conn.commit()
    conn.close()
    print("Data inserted successfully!")

# خواندن داده بر اساس id
def get_data_by_id(record_id):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM speeches WHERE id = ?', (record_id,))
    result = cursor.fetchone()

    conn.close()
    return result

# خواندن داده بر اساس part_id
def get_data_by_part_id(part_id):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM speeches WHERE part_id = ?', (part_id,))
    result = cursor.fetchall()

    conn.close()
    return result

# خواندن داده بر اساس context_id
def get_data_by_context_id(context_id):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM speeches WHERE context_id = ?', (context_id,))
    result = cursor.fetchall()

    conn.close()
    return result

# خواندن تمام داده‌ها
def get_all_data():
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM speeches')
    result = cursor.fetchall()

    conn.close()
    return result

# متد ایجاد جدول چت
def create_chat_table():
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    # ایجاد جدول اگر وجود نداشته باشد
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS chat (
        id TEXT PRIMARY KEY,
        title TEXT,
        user_id TEXT,
        user_query TEXT,
        model_key TEXT,
        retrived_passage TEXT,
        retrived_ref_ids TEXT,
        prompt_type TEXT,
        retrived_duration TEXT,
        llm_duration TEXT,
        full_duration TEXT,
        time_create TEXT,
        used_ref_ids TEXT,
        prompt_answer TEXT,
        status BOOLEAN
    )
    ''')
    conn.commit()
    conn.close()
    print("Chat table created successfully!")

# متد ذخیره یک داده تکی در جدول چت
def insert_chat_message(chat_obj):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    # تبدیل لیست‌ها یا دیکشنری‌ها به JSON برای ذخیره در جدول
    chat_obj['retrived_ref_ids'] = json.dumps(chat_obj['retrived_ref_ids'])
    chat_obj['used_ref_ids'] = json.dumps(chat_obj['used_ref_ids'])

    try:
        cursor.execute('''
        INSERT INTO chat (id, title, user_id, user_query, model_key, retrived_passage, retrived_ref_ids, prompt_type, retrived_duration, llm_duration, full_duration, time_create, used_ref_ids, prompt_answer, status)
        VALUES (:id, :title, :user_id, :user_query, :model_key, :retrived_passage, :retrived_ref_ids, :prompt_type, :retrived_duration, :llm_duration, :full_duration, :time_create, :used_ref_ids, :prompt_answer, :status)
        ''', chat_obj)
        conn.commit()
        print(f"Chat message with id '{chat_obj['id']}' inserted successfully!")
    except sqlite3.IntegrityError:
        print(f"Warning: Chat message with id '{chat_obj['id']}' already exists!")
    conn.close()

# متد بازیابی داده بر اساس آی‌دی
def get_chat_message_by_id(chat_id):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM chat WHERE id = ?', (chat_id,))
    result = cursor.fetchone()

    conn.close()

    if result:
        # تبدیل ستون‌های JSON به نوع اصلی (لیست یا دیکشنری)
        result_dict = {
            "id": result[0],
            "title": result[1],
            "user_id": result[2],
            "user_query": result[3],
            "model_key": result[4],
            "retrived_passage": result[5],
            "retrived_ref_ids": json.loads(result[6]),
            "prompt_type": result[7],
            "retrived_duration": result[8],
            "llm_duration": result[9],
            "full_duration": result[10],
            "time_create": result[11],
            "used_ref_ids": json.loads(result[12]),
            "prompt_answer": result[13],
            "status": bool(result[14])
        }
        return result_dict
    return None

# متد بازیابی چت‌ها بر اساس user_id
def get_chats_by_user_id(user_id):
    conn = sqlite3.connect('./db/nahj.db')
    cursor = conn.cursor()

    cursor.execute('SELECT * FROM chat WHERE user_id = ?', (user_id,))
    results = cursor.fetchall()

    conn.close()

    # تبدیل هر رکورد به دیکشنری و مدیریت ستون‌های JSON
    chats = []
    for result in results:
        chat = {
            "id": result[0],
            "title": result[1],
            "user_id": result[2],
            "user_query": result[3],
            "model_key": result[4],
            "retrived_passage": result[5],
            "retrived_ref_ids": json.loads(result[6]),
            "prompt_type": result[7],
            "retrived_duration": result[8],
            "llm_duration": result[9],
            "full_duration": result[10],
            "time_create": result[11],
            "used_ref_ids": json.loads(result[12]),
            "prompt_answer": result[13],
            "status": bool(result[14])
        }
        chats.append(chat)

    return chats

def create_tables():
    create_speechs_table()
    create_chat_table()

# مثال استفاده
if __name__ == "__main__":
    # ایجاد جدول (فقط بار اول اجرا می‌شود)
    #create_tables()


    # درج داده‌ها
    with open('./data-faiss/faiss_index_nahj_metadata.json', 'r', encoding='utf-8') as file:
        data = json.load(file)
        
    insert_data(data)

    # خواندن داده بر اساس id
    record = get_data_by_id("wisdom774")
    print("Record by ID:", record)

    # خواندن داده بر اساس part_id
    records_by_part = get_data_by_part_id("sp13")
    print("Records by Part ID:", records_by_part)

    # خواندن داده بر اساس context_id
    records_by_context = get_data_by_context_id("sn1")
    print("Records by Context ID:", records_by_context)