import json
import os
import re
import numpy as np
import faiss
from langchain_openai import OpenAIEmbeddings
from dotenv import load_dotenv

# Load env variables for OpenAI
load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".env"))

def main():
    # 1. Trỏ thẳng vào thư mục datadb
    current_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(current_dir, "tonghop.txt")

    if not os.path.exists(file_path):
        print(f" File không tồn tại: {file_path}")
        return

    with open(file_path, encoding="utf-8") as f:
        content = f.read()

    # 2. Bóc tách file theo các đầu mục "FILE: ..."
    sections = re.split(r"={20,}\nFILE:\s*(.*?)\n={20,}\n", content)

    chunks = []
    if len(sections) > 1:
        # Bỏ qua index 0
        for i in range(1, len(sections), 2):
            title = sections[i].strip()
            text = sections[i + 1].strip()
            if text:
                chunks.append({"id": len(chunks), "title": title, "content": text})
    else:
        print(" Không tìm thấy delimiter 'FILE:' trong tonghop.txt.")
        return

    print(f" BÓC TÁCH: Đã chia thành {len(chunks)} phần chính sách/tài liệu.")

    # 3. Embedding với OpenAI
    print("\n[Loading Model] Khởi tạo OpenAIEmbeddings (text-embedding-3-small)...")
    embedder = OpenAIEmbeddings(model="text-embedding-3-small")

    # Gom title + content để embed
    embed_texts = [f"{c['title']}\n{c['content']}" for c in chunks]

    print("[Embedding] Đang gọi API OpenAI để sinh Vectors...")
    embeddings_list = embedder.embed_documents(embed_texts)
    embeddings = np.array(embeddings_list).astype("float32")

    # Chuẩn hoá L2 để tính Cosine Similarity mượt hơn với FAISS IndexFlatIP
    faiss.normalize_L2(embeddings)

    # 4. Khởi tạo FAISS Index & Lưu metadata
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)

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

    faiss.write_index(index, index_path)
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(chunks, f, ensure_ascii=False, indent=2)

    print(f" ĐÃ Lưu FAISS vector vào: {index_path} (dim={dim})")
    print(f" ĐÃ Lưu JSON nội dung vào: {meta_path}\n")

    # 5. TEST MÔ PHỎNG LUÔN
    queries = [
        "mua hàng online muốn đổi trực tiếp ở cửa hàng gần nhất thì phải làm sao hả shop?",
        "thời gian gửi hàng ship về hà đông mất tầm bao lâu?",
        "có bán quần của trẻ em bé gái cao tầm 1 mét 4 không cho xin thông số",
    ]

    print(" BẮT ĐẦU TEST SEARCH RAG (TRUY VẤN FULL CÒN NGUYÊN BẢN VỚI OPENAI EMBEDDING)")
    print("=" * 70)

    for q in queries:
        print(f"\n CÂU HỎI: '{q}'")

        # Sinh vector câu hỏi bằng OpenAI
        q_emb_list = embedder.embed_query(q)
        q_emb = np.array([q_emb_list]).astype("float32")
        faiss.normalize_L2(q_emb)

        # Search top 1
        scores, I = index.search(q_emb, 1)

        best_id = I[0][0]
        best_score = scores[0][0]

        if best_id != -1:
            best_chunk = chunks[best_id]
            print(f" MATCH ĐẦU MỤC: '{best_chunk['title']}' (Score: {best_score:.4f})")
            print(
                f" NỘI DUNG FULL RETURN (Trích 300 chữ): {best_chunk['content'][:300].replace(chr(10), ' ')} ... <CÒN NỮA>"
            )
            print(f"   Độ dài thực tế sẽ đưa cho LLM: {len(best_chunk['content'])} ký tự.")
        else:
            print("Không tìm thấy.")

if __name__ == "__main__":
    main()
