import os
import pymysql
import uuid
from flask import Flask, render_template, request, jsonify, session
import markdown

#from llama_index.llms.gemini import Gemini
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import TokenTextSplitter
from llama_index.core import StorageContext, load_index_from_storage
from llama_index.core import PromptTemplate
import hashlib
from bs4 import BeautifulSoup
from flask_cors import CORS

#from langchain.text_splitter import RecursiveCharacterTextSplitter
# Khởi tạo Flask app
app = Flask(__name__)
CORS(app)  # This allows all origins by default
app.secret_key = os.urandom(24)

# Thiết lập API key và các biến môi trường
google_gemini_api = "AIzaSyAv_HWvX7agGiktA99HqDQTa-tX5vtZDRA"
os.environ["GOOGLE_API_KEY"] = google_gemini_api
os.environ["MODEL_NAME"] = "models/gemini-2.5-flash"

# Cấu hình mô hình
#llm = Gemini(model_name="models/gemini-2.0-flash-001", api_key=os.environ["GOOGLE_API_KEY"])
llm = GoogleGenAI(model="models/gemini-2.5-flash", api_key=os.environ["GOOGLE_API_KEY"])
#embed_model = HuggingFaceEmbedding(model_name="Viet-Mistral/Vistral-7B-Chat")
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-m3")
Settings.llm = llm
Settings.embed_model = embed_model


    
# Hàm tính hash của file để kiểm tra thay đổi
def get_file_hash(file_path):
    with open(file_path, "rb") as f:
        content = f.read()
        return hashlib.md5(content).hexdigest()

# Lấy danh sách file và hash từ thư mục (bao gồm subfolder)
def get_all_files_info(data_dir):
    files_info = {}
    for root, _, files in os.walk(data_dir):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            # Tính hash cho file
            file_hash = get_file_hash(file_path)
            # Lưu đường dẫn tương đối để dễ quản lý
            relative_path = os.path.relpath(file_path, data_dir)
            files_info[relative_path] = file_hash
    return files_info

# Lưu thông tin file đã xử lý
def load_processed_files_info(persist_dir="ctump_index_storage"):
    info_file = os.path.join(persist_dir, "processed_files.json")
    if os.path.exists(info_file):
        with open(info_file, "r") as f:
            import json
            return json.load(f)
    return {}

def save_processed_files_info(processed_files, persist_dir="ctump_index_storage"):
    info_file = os.path.join(persist_dir, "processed_files.json")
    with open(info_file, "w") as f:
        import json
        json.dump(processed_files, f)

# Xóa node cũ liên quan đến file
def remove_old_nodes(index, file_path):
    # Duyệt qua tất cả node trong index
    nodes_to_remove = []
    for node_id, node in index.docstore.docs.items():
        # Kiểm tra metadata của node
        if node.metadata.get("file_path") == file_path:
            nodes_to_remove.append(node_id)
    
    # Xóa các node cũ
    for node_id in nodes_to_remove:
        index.delete_nodes([node_id])
    print(f"Removed {len(nodes_to_remove)} old nodes for {file_path}")

# Tạo hoặc cập nhật index
def create_or_update_index():
    PERSIST_DIR = "ctump_index_storage"
    DATA_DIR = "data_ctump"

    
    
    # Load thông tin file đã xử lý
    processed_files = load_processed_files_info(PERSIST_DIR)
    current_files = get_all_files_info(DATA_DIR)
    
    # Kiểm tra file mới hoặc thay đổi
    files_to_process = [f for f, hash_val in current_files.items() 
                       if f not in processed_files or processed_files[f] != hash_val]
    
    if not os.path.exists(PERSIST_DIR) or not processed_files:
        # Nếu chưa có index, tạo mới toàn bộ
        reader = SimpleDirectoryReader(input_dir=DATA_DIR, recursive=True)
        documents = reader.load_data()
        splitter = TokenTextSplitter(chunk_size=1024, chunk_overlap=100, separator=".\n")
        nodes = splitter.get_nodes_from_documents(documents)
        index = VectorStoreIndex(nodes)
        index.storage_context.persist(persist_dir=PERSIST_DIR)
    elif files_to_process:
        # Chỉ xử lý file mới hoặc thay đổi
        reader = SimpleDirectoryReader(input_files=[os.path.join(DATA_DIR, f) for f in files_to_process], recursive=True)
        documents = reader.load_data()
        splitter = TokenTextSplitter(chunk_size=1024, chunk_overlap=100, separator=".\n")
        nodes = splitter.get_nodes_from_documents(documents)

        # Load index cũ và thêm nodes mới
        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
        index = load_index_from_storage(storage_context)
        for node in nodes:
            index.insert_nodes([node])  # Thêm từng node mới
        index.storage_context.persist(persist_dir=PERSIST_DIR)
    else:
        # Không có thay đổi, load index cũ
        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
        index = load_index_from_storage(storage_context)
            
    # Cập nhật thông tin file đã xử lý
    processed_files.update(current_files)
    save_processed_files_info(processed_files, PERSIST_DIR)
    return index

# Sử dụng index
index = create_or_update_index()

# Thiết lập Prompt
# QA_PROMPT_TMPL = (
#     "Không trả lời các câu hỏi về chính trị, bạo lực. Sử dụng các phần ngữ cảnh sau để trả lời câu hỏi của người dùng. Bạn tên là chatbot CTUMP của trường đại học y dược Cần Thơ. Cố gắng tìm tài liệu đã truy vấn trả lời ra"
#     "Nếu không biết câu trả lời, bạn chỉ cần nói rằng bạn không biết. Bạn phải trả lời bằng ngôn ngữ của câu hỏi.\n"
#     "---------------------\n"
#     "{context_str}\n"
#     "---------------------\n"
#     "Câu hỏi: {query_str}\n"
#     "Trả lời: "
# )
QA_PROMPT_TMPL = (
    "Không trả lời các câu hỏi về chính trị, bạo lực.Không dẫn link từ tài liệu huấn luyện. Sử dụng các phần ngữ cảnh sau để trả lời câu hỏi của người dùng. Cố gắng tìm tài liệu đã truy vấn trả lời ra"
    "Nếu không biết câu trả lời, bạn có thể tổng hợp thông tin từ internet để trả lời. Bạn phải trả lời bằng ngôn ngữ của câu hỏi."
    "Sử dụng tiêu đề, gạch đầu dòng để tổ chức thông tin. In đậm các từ khóa quan trọng để nhấn mạnh."
    "Phân tích kỹ lưỡng dữ liệu bảng được cung cấp, bao gồm các cột, hàng và giá trị, để hiểu rõ ý nghĩa và mối quan hệ giữa các dữ liệu."
    "Chuyển đổi thông tin từ bảng thành một đoạn văn liên kết, sử dụng ngôn ngữ tự nhiên, rõ ràng và phù hợp với ngữ cảnh, tuyệt đối không sử dụng định dạng bảng hoặc liệt kê kiểu bảng trong câu trả lời."
    "Nếu ngữ cảnh không rõ, hãy đưa ra giả định hợp lý và đề xuất câu hỏi dựa trên giả định đó, đồng thời cho phép người dùng làm rõ nếu cần."
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "Câu hỏi: {query_str}\n"
    "Trả lời: "
)
qa_prompt = PromptTemplate(QA_PROMPT_TMPL)
query_engine = index.as_query_engine(
    similarity_top_k=10,  # Tăng số lượng kết quả trả về
    max_marginal_relevance=True,  # Sử dụng MMR để chọn kết quả có độ tương đồng cao và đa dạng hơn
    alpha=1  # Điều chỉnh độ kết hợp giữa sự tương đồng và sự đa dạng
)
query_engine.update_prompts({"response_synthesizer:text_qa_template": qa_prompt})

def get_db_connection():
    try:
        return pymysql.connect(
            host="localhost",
            user="root",
            #password="",
            password="Pasw@rd1473",
            database="chatbot"
        )
    except pymysql.Error as err:
        print(f"Không thể kết nối MySQL: {err}")
        return None

# Hàm lưu vào database với da_duoc_hoc là chuỗi "True" hoặc "False"
def save_to_db(cauhoi, traloi, session_id):
    conn = get_db_connection()
    
    if conn is None:
        print("Không thể lưu vào database do lỗi kết nối.")
        return
    cursor = conn.cursor()
        
    traloi = traloi.replace("<br>", "")
    query = "INSERT INTO lichsuchat (session_id, cauhoi, traloi, thoigian, da_duoc_hoc) VALUES (%s, %s, %s, CURRENT_TIMESTAMP, 0)"
    values = (session_id, cauhoi, traloi)
    cursor.execute(query, values)
    conn.commit()
    cursor.close() 
    conn.close() 
    

# Trang chính hiển thị giao diện
@app.route("/")
def index():
    return render_template("index.html")


@app.route("/check_spell", methods=["POST"])
def check_spell():
    data = request.get_json()
    user_question = data.get("question")  # Sử dụng trường 'question'

    return jsonify({"response": user_question})  # Sử dụng trường 'response'
    
    
# API để hỏi câu hỏi
@app.route("/ask", methods=["POST"])
def ask():
    data = request.get_json()
    user_question = data.get("question")
    print(f"Tin nhắn từ người dùng: {user_question}")

    if 'session_id' not in session:
        session['session_id'] = str(uuid.uuid4())
    session_id = session['session_id']
    # Proceed with database logic

    try:
        response = query_engine.query(user_question)
        html = markdown.markdown(str(response), extensions=['tables'])
        # Sử dụng BeautifulSoup để thêm target="_blank"
        soup = BeautifulSoup(html, 'html.parser')
        for a_tag in soup.find_all('a'):
            a_tag['target'] = '_blank'
            a_tag['rel'] = 'noopener noreferrer'  # Bảo mật
        response = str(soup)
        # Lưu vào database với giá trị da_duoc_hoc phù hợp
        save_to_db(user_question, response, session_id)
        return jsonify({"response": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Khởi chạy server Flask
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)