import json
import os

import faiss
import numpy as np
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

current_dir = os.path.dirname(os.path.abspath(__file__))
load_dotenv(os.path.join(current_dir, "..", ".env"))


def main():
    print(" BẮT ĐẦU TEST E2E RAG VỚI GPT-4o-mini (BẢN PROMPT THÔNG MINH HƠN)")
    print("=" * 80)

    index_path = os.path.join(current_dir, "canifa_docs.index")
    meta_path = os.path.join(current_dir, "canifa_docs_meta.json")

    if not os.path.exists(index_path) or not os.path.exists(meta_path):
        print(" Chưa có Index. Chạy file build_and_test_faiss.py trước để tạo.")
        return

    print("[1] Đang nạp thư viện Vector từ FAISS và metadata...")
    index = faiss.read_index(index_path)
    with open(meta_path, encoding="utf-8") as f:
        chunks = json.load(f)

    print("[2] Khởi tạo Embedder và Models (gpt-4o-mini)...")
    embedder = OpenAIEmbeddings(model="text-embedding-3-small")
    llm = ChatOpenAI(model="gpt-4.1-mini", temperature=0)

    prompt_template = PromptTemplate.from_template(
        "Bạn là nhân viên CSKH của Canifa. Bạn luôn nhiệt tình và vui vẻ.\n"
        "Yêu cầu:\n"
        "1. Nếu khách hàng chỉ chào hỏi (ví dụ: 'chào', 'hello'), hãy chào lại lịch sự và hỏi xem bạn có thể giúp gì được không, KHÔNG báo lỗi thiếu thông tin.\n"
        "2. Nếu khách hỏi về chính sách, quy định, cửa hàng, hãy DỰA VÀO TÀI LIỆU CUNG CẤP BÊN DƯỚI để trả lời chi tiết và thân thiện.\n"
        "3. Đọc thật kỹ tài liệu được cung cấp. Nếu tài liệu ĐÃ CÓ nhắc đến từ khoá khách hỏi, hãy cố gắng tóm tắt câu trả lời thay vì từ chối.\n"
        "4. CHỈ KHI NÀO tài liệu dưới đây hoàn toàn không nhắc đến nội dung khách hỏi, bạn mới được nói: 'Dạ em chưa tìm thấy thông tin này trong chính sách hiện tại. Bạn vui lòng liên hệ hotline 1800 6061 nha!'\n"
        "5. Không bao giờ được tự bịa ra thông số, giá tiền, hoặc điều khoản nếu tài liệu không có.\n\n"
        "--- BẮT ĐẦU TÀI LIỆU ---\n"
        "{context}\n"
        "--- KẾT THÚC TÀI LIỆU ---\n\n"
        "CÂU HỎI CỦA KHÁCH: {question}\n\n"
        "TRẢ LỜI CỦA BẠN:"
    )

    chain = prompt_template | llm

    print("\n" + "*" * 80)
    print(" CHATBOT ĐÃ SẴN SÀNG! (Gõ 'exit' hoặc 'quit' để thoát)")
    print("*" * 80 + "\n")

    while True:
        try:
            q = input("\n BẠN: ")
            if q.strip().lower() in ["exit", "quit"]:
                print(" Tạm biệt!")
                break

            if not q.strip():
                continue

            q_emb_list = embedder.embed_query(q)
            q_emb = np.array([q_emb_list]).astype("float32")
            faiss.normalize_L2(q_emb)

            TOP_K = 10
            scores, I = index.search(q_emb, TOP_K)

            contexts = []
            for i in range(TOP_K):
                doc_id = I[0][i]
                if doc_id != -1:
                    doc_title = chunks[doc_id]["title"]
                    doc_content = chunks[doc_id]["content"]
                    contexts.append(f"--- TÀI LIỆU: {doc_title} ---\n{doc_content}")

            if contexts:
                combined_context = "\n\n".join(contexts)
                response = chain.invoke({"context": combined_context, "question": q})
                print(f" CANIFA AI:\n{response.content}")
            else:
                print(" CANIFA AI: Hệ thống không tìm thấy tài liệu phù hợp.")

        except (KeyboardInterrupt, EOFError):
            print("\n Tạm biệt!")
            break


if __name__ == "__main__":
    main()
