backend/app/services/chat_service.py

"""
DeepMed-AI — services/chat_service.py
HybridRAG ChatService: ChromaDB + BM25 + CrossEncoder Reranker.

Architecture (inspired by reference chatbot, upgraded for FastAPI):
  - Embedding: paraphrase-multilingual-MiniLM-L12-v2
  - Vector DB: ChromaDB (local persistent)
  - Hybrid: BM25 keyword + Vector semantic (EnsembleRetriever 50/50)
  - Reranker: BGE-reranker-v2-m3 (Deep mode only)
  - 2 modes: Fast (Ensemble k=15) vs Deep (Ensemble k=25 → Reranker top 5)
  - LLM: Gemini 2.5 Flash via LangChain
  - History-aware retrieval via create_history_aware_retriever
"""

import os
import re
import shutil
import logging
import traceback
import unicodedata
from datetime import datetime
from typing import Any, Dict, List, Optional

import pandas as pd

from langchain_core.documents import Document
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain, create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain

from app.core.config import GOOGLE_API_KEY, GOOGLE_API_KEYS, NVIDIA_API_KEY, OPENROUTER_API_KEY, DATA_DIR, CHROMA_DB_PATH
from app.core.logging_config import logger
from app.tools.llm_client import get_llm, get_llm_name
from app.services.database_service import db_service

# ── Constants ──────────────────────────────────────────────────────────────────
MAX_HISTORY_TURNS = 5
FORCE_REBUILD_DB = os.getenv("FORCE_REBUILD_DB", "False").lower() == "true"

# Detect device for model inference
try:
    import torch
    DEVICE = "cuda" if torch.cuda.is_available() else (
        "mps" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else "cpu"
    )
except ImportError:
    DEVICE = "cpu"


# ── System Prompt (giữ nguyên) ─────────────────────────────────────────────────
DEEPMED_SYSTEM_PROMPT = """Bạn là **DeepMed-AI** — Trợ lý Y khoa và Dược lâm sàng của **Trung tâm Y tế Khu vực Thanh Ba**, \
tỉnh Phú Thọ, Việt Nam. Bạn hỗ trợ đội ngũ y bác sĩ và dược sĩ trong:
- Tra cứu thông tin thuốc (tên, hoạt chất, **giá**, hãng sản xuất, liều dùng, chống chỉ định...)
- Chẩn đoán, phác đồ điều trị, tư vấn y khoa
- Tra cứu danh mục thuốc nội bộ của TTYT Thanh Ba

## Nguyên tắc bắt buộc:

1. **Chính xác & Trung thực**: Chỉ trả lời dựa trên tài liệu Context. \
Nếu Context không có thông tin, hãy nói rõ: "Trong kho dữ liệu hiện tại chưa có thông tin về vấn đề này", \
tuyệt đối không tự bịa phác đồ hay thông tin thuốc.

2. **Giá thuốc**: Khi tài liệu nội bộ có thông tin giá (ví dụ "Giá của thuốc X: 21.798"), \
hãy TRẢ LỜI giá đó. KHÔNG được nói "tôi không cung cấp thông tin giá". \
Đây là thông tin danh mục thuốc nội bộ, KHÔNG phải thông tin thương mại.

3. **Trích dẫn đúng nguồn**: \
   - Chỉ sử dụng thông tin từ các tài liệu Context THỰC SỰ LIÊN QUAN đến câu hỏi.\
   - Nếu tài liệu không liên quan đến nội dung câu trả lời → BỎ QUA, KHÔNG trích dẫn.\
   - KHÔNG bịa tên tài liệu.

4. **Văn phong**: Hãy trình bày như một *đồng nghiệp* đang trao đổi chuyên môn: \
tự nhiên, trôi chảy, dễ hiểu. In đậm từ khóa quan trọng, dùng gạch đầu dòng khi cần. \
Trả lời bằng tiếng Việt, kèm thuật ngữ Latin/Anh trong ngoặc nếu cần.

5. **Quy tắc an toàn**:
   - KHÔNG BAO GIỜ lấy thông tin thuốc A gán cho thuốc B
   - Không chẩn đoán thay bác sĩ — chỉ hỗ trợ tra cứu
   - Cấp cứu: hướng dẫn sơ cứu + gọi 115

6. **Xác định thuốc đúng**: Mỗi tài liệu bắt đầu bằng header `[Trích từ tài liệu: TÊN_FILE]`. \
Dùng header này để xác định tài liệu thuộc về thuốc nào. Nếu câu hỏi về thuốc X mà tài liệu \
là về thuốc Y → BỎ QUA tài liệu đó, không sử dụng.

7. **Bối cảnh**: Ưu tiên phác đồ Bộ Y tế Việt Nam, WHO, thực tế tuyến huyện.

8. **Luôn kết thúc bằng**: *Thông tin mang tính chất tham khảo từ dữ liệu nội bộ. \
Quyết định điều trị thuộc về bác sĩ chuyên khoa.*

9. **Khai báo nguồn đã dùng (BẮT BUỘC)**: Ở cuối cùng câu trả lời, \
hãy liệt kê chính xác tên file tài liệu bạn ĐÃ THỰC SỰ SỬ DỤNG để sinh ra câu trả lời, \
theo format: `[USED_SOURCES: tên_file_1 | tên_file_2 | ...]`. \
CHỈ liệt kê tài liệu bạn thực sự trích dẫn/sử dụng, KHÔNG liệt kê tài liệu không liên quan. \
Nếu không dùng tài liệu nào, ghi `[USED_SOURCES: NONE]`."""

# Deep mode prompt — 5-step clinical reasoning with drug matching
QA_SYSTEM_PROMPT_DEEP = (
    DEEPMED_SYSTEM_PROMPT + "\n\n"
    "🧠 QUY TRÌNH TƯ DUY LÂM SÀNG (Suy luận nội bộ — KHÔNG in tên bước, chỉ thể hiện kết quả):\n\n"
    "**Bước 1 — Phác đồ & Hướng dẫn điều trị:**\n"
    "Quét Context tìm các Phác đồ Bộ Y tế (QĐ-BYT), Hướng dẫn của các Hiệp hội (VUNA, GOLD, GINA, IDSA...), "
    "và khuyến cáo WHO liên quan đến bệnh lý được hỏi. Trích xuất: tiêu chuẩn chẩn đoán, phân loại mức độ, "
    "phác đồ điều trị bậc thang/lựa chọn đầu tay, các thuốc/nhóm thuốc khuyến cáo kèm liều dùng.\n\n"
    "**Bước 2 — Cảnh giác dược & An toàn thuốc:**\n"
    "Kiểm tra Context có cảnh báo cảnh giác dược, ADR, tương tác thuốc, thu hồi thuốc, "
    "hoặc lưu ý an toàn nào liên quan đến bệnh lý hoặc nhóm thuốc ở Bước 1 hay không. "
    "Ghi nhận các chống chỉ định đặc biệt, cặp tương tác nguy hiểm, hiệu chỉnh liều trên đối tượng đặc biệt "
    "(suy thận, suy gan, thai kỳ, người cao tuổi).\n\n"
    "**Bước 3 — Đối chiếu thuốc nội bộ TTYT Thanh Ba:**\n"
    "Từ danh sách thuốc/hoạt chất khuyến cáo ở Bước 1, đối chiếu với Danh mục thuốc nội bộ trong Context "
    "(file DANH_MUC_THUOC_NOI_BO_TOAN_BO hoặc các chunk chứa thông tin thuốc nội bộ). "
    "Xác định: thuốc nào CÓ SẴN tại kho, tên biệt dược, hoạt chất, hàm lượng, đường dùng, giá. "
    "Nếu thuốc khuyến cáo KHÔNG có sẵn → đề xuất thuốc thay thế cùng nhóm/cùng hoạt chất có trong kho.\n\n"
    "**Bước 4 — Tổng hợp phác đồ điều trị thực tế:**\n"
    "Đúc kết thành phác đồ điều trị khả thi tại TTYT Thanh Ba, bao gồm: "
    "(a) Nhận định lâm sàng tóm tắt, (b) Phác đồ thuốc cụ thể với liều dùng — "
    "ưu tiên thuốc có sẵn nội bộ, (c) Lưu ý cảnh giác dược/tương tác/ADR quan trọng, "
    "(d) Tiêu chí theo dõi và đánh giá đáp ứng.\n\n"
    "**Bước 5 — Trình bày:**\n"
    "Viết câu trả lời MỘT CÁCH TỰ NHIÊN, MẠCH LẠC theo dàn ý sau (có thể linh hoạt):\n"
    "  1. **Tổng quan** — Tóm tắt hướng dẫn điều trị theo phác đồ/hiệp hội (trích nguồn).\n"
    "  2. **Phác đồ điều trị đề xuất** — Liệt kê thuốc cụ thể CÓ SẴN tại TTYT Thanh Ba "
    "(tên thuốc, hoạt chất, liều, đường dùng). Nếu không có thuốc phù hợp, nêu rõ.\n"
    "  3. **Lưu ý an toàn** — Cảnh giác dược, tương tác thuốc, chống chỉ định đặc biệt.\n"
    "  4. **Theo dõi** — Tiêu chí đánh giá đáp ứng, thời điểm tái khám (nếu có trong Context).\n\n"
    "🎯 NGUYÊN TẮC TRÌNH BÀY:\n"
    "- Nếu nhắc đến thuốc hoặc phác đồ, nói rõ nó thuộc tài liệu nào trong Context.\n"
    "- Khi đề xuất thuốc nội bộ, ghi rõ: **tên biệt dược** (hoạt chất, hàm lượng) — giá nếu có.\n"
    "- Trình bày đẹp mắt: heading, bullet points, bảng nếu cần.\n"
    "- KHÔNG tự bịa thuốc không có trong Context.\n\n"
    "*Lưu ý: Phải chốt lại bằng dòng in nghiêng: "
    "'Thông tin mang tính chất tham khảo từ dữ liệu nội bộ. Quyết định điều trị thuộc về bác sĩ chuyên khoa.'*\n\n"
    "DỮ LIỆU NỘI BỘ (Context):\n{context}"
)

# Fast mode prompt — direct answer extraction
QA_SYSTEM_PROMPT_FAST = (
    DEEPMED_SYSTEM_PROMPT + "\n\n"
    "DỮ LIỆU NỘI BỘ (Context):\n{context}"
)


# ══════════════════════════════════════════════════════════════════════════════
#  Document Loading (multi-format: PDF, DOCX, XLSX, CSV, TXT, MD)
# ══════════════════════════════════════════════════════════════════════════════

def _process_excel_file(file_path: str, filename: str) -> List[Document]:
    """Each Excel/CSV row becomes a separate Document."""
    docs = []
    try:
        if file_path.endswith(".csv"):
            df = pd.read_csv(file_path)
        else:
            df = pd.read_excel(file_path)
        df.dropna(how="all", inplace=True)
        df.fillna("Không có thông tin", inplace=True)
        for idx, row in df.iterrows():
            content_parts = []
            for col_name, val in row.items():
                clean_val = str(val).strip()
                if clean_val and clean_val.lower() != "nan":
                    content_parts.append(f"{col_name}: {clean_val}")
            if content_parts:
                page_content = "\n".join(content_parts)
                metadata = {
                    "source": filename,
                    "source_path": file_path,
                    "row": idx + 1,
                    "type": "excel_record",
                }
                docs.append(Document(page_content=page_content, metadata=metadata))
    except Exception as e:
        logger.error("Lỗi xử lý Excel %s: %s", filename, e)
    return docs


def _load_documents_from_folder(folder_path: str) -> List[Document]:
    """Recursively load all supported files from the data directory."""
    logger.info("--- Bắt đầu quét thư mục: %s ---", folder_path)
    documents: List[Document] = []
    if not os.path.exists(folder_path):
        os.makedirs(folder_path, exist_ok=True)
        return []

    for root, dirs, files in os.walk(folder_path):
        # Skip individual drug .md files — all info is in DANH_MUC_THUOC_NOI_BO_TOAN_BO.md
        root_norm = root.replace("\\", "/").lower()
        if "thông tin thuốc nội bộ" in root_norm or "thong tin thuoc noi bo" in root_norm:
            continue
        for filename in files:
            file_path = os.path.join(root, filename)
            filename_lower = filename.lower()
            try:
                if filename_lower.endswith(".pdf"):
                    from langchain_community.document_loaders import PyPDFLoader
                    loader = PyPDFLoader(file_path)
                    docs = loader.load()
                    for d in docs:
                        d.metadata["source"] = filename
                        d.metadata["source_path"] = file_path
                    documents.extend(docs)

                elif filename_lower.endswith(".docx"):
                    try:
                        import docx2txt
                        text = docx2txt.process(file_path)
                    except ImportError:
                        import docx
                        doc = docx.Document(file_path)
                        text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
                    if text and text.strip():
                        documents.append(Document(
                            page_content=text,
                            metadata={"source": filename, "source_path": file_path},
                        ))

                elif filename_lower.endswith((".xlsx", ".xls", ".csv")):
                    excel_docs = _process_excel_file(file_path, filename)
                    documents.extend(excel_docs)

                elif filename_lower.endswith((".txt", ".md")):
                    with open(file_path, "r", encoding="utf-8") as f:
                        text = f.read()
                    if text.strip():
                        documents.append(Document(
                            page_content=text,
                            metadata={"source": filename, "source_path": file_path},
                        ))

            except Exception as e:
                logger.error("Lỗi đọc file %s: %s", filename, e)

    logger.info("Tổng cộng đã load: %d tài liệu gốc.", len(documents))
    return documents


def _normalize_for_search(text: str) -> str:
    """Normalize Vietnamese text for accent-insensitive keyword matching."""
    if not text:
        return ""
    normalized = unicodedata.normalize("NFD", text)
    normalized = "".join(ch for ch in normalized if unicodedata.category(ch) != "Mn")
    normalized = normalized.lower()
    normalized = re.sub(r"[^a-z0-9\s+-]", " ", normalized)
    return re.sub(r"\s+", " ", normalized).strip()


def _extract_drug_profile(doc: Document) -> Dict[str, Any]:
    """Extract drug-specific metadata from internal drug markdown files."""
    source = str(doc.metadata.get("source", ""))
    source_path = str(doc.metadata.get("source_path", ""))
    source_norm = (source_path or source).replace("\\", "/").lower()
    content = doc.page_content or ""

    if "thông tin thuốc nội bộ" not in source_norm and "thong tin thuoc noi bo" not in source_norm:
        return {}

    # Drug name from heading: # TÊN_THUỐC
    m_name = re.search(r"^\s*#\s+(.+)$", content, re.MULTILINE)
    drug_name = m_name.group(1).strip() if m_name else os.path.splitext(os.path.basename(source))[0]

    # Active ingredient from line: Hoạt chất: ...
    m_active = re.search(r"^\s*Hoạt\s*chất\s*:\s*(.+)$", content, re.IGNORECASE | re.MULTILINE)
    active_ingredient = m_active.group(1).strip() if m_active else ""

    # Ingredient keyword list for lexical retrieval reinforcement
    ingredient_tokens = []
    if active_ingredient:
        base = _normalize_for_search(active_ingredient)
        tokens = re.split(r"[\s+,;()/-]+", base)
        stop = {
            "mg", "ml", "mcg", "ui", "iu", "g", "kg", "hoat", "chat", "duoi", "dang",
            "vien", "ong", "goi", "tiem", "uong", "acid",
        }
        ingredient_tokens = sorted({t for t in tokens if len(t) >= 3 and t not in stop and not t.isdigit()})

    return {
        "doc_type": "drug_info",
        "drug_name": drug_name,
        "drug_name_upper": drug_name.upper(),
        "active_ingredient": active_ingredient,
        "active_ingredient_norm": _normalize_for_search(active_ingredient),
        "ingredient_keywords": ", ".join(ingredient_tokens),
    }


def _bm25_preprocess(text: str) -> List[str]:
    """Lowercase + simple tokenization for BM25 so matching is case-insensitive."""
    return text.lower().split()


def _extract_drug_profile_from_content(content: str) -> Dict[str, Any]:
    """Extract drug profile from any text content (not requiring specific folder)."""
    if not content or len(content) < 50:
        return {}

    # Drug name from heading: # TÊN_THUỐC
    m_name = re.search(r"^\s*#\s+(.+)$", content, re.MULTILINE)
    drug_name = m_name.group(1).strip() if m_name else ""
    if not drug_name:
        return {}

    # Active ingredient from line: Hoạt chất: ...
    m_active = re.search(r"^\s*Hoạt\s*chất\s*:\s*(.+)$", content, re.IGNORECASE | re.MULTILINE)
    active_ingredient = m_active.group(1).strip() if m_active else ""

    # Price from line: Giá: ... or **[TRA CỨU NHANH]** ... Giá: ...
    price = ""
    m_price = re.search(r"Giá\s*:\s*([\d.,]+)", content)
    if m_price:
        price = m_price.group(1).strip()

    # Ingredient keyword list
    ingredient_tokens = []
    if active_ingredient:
        base = _normalize_for_search(active_ingredient)
        tokens = re.split(r"[\s+,;()/-]+", base)
        stop = {
            "mg", "ml", "mcg", "ui", "iu", "g", "kg", "hoat", "chat", "duoi", "dang",
            "vien", "ong", "goi", "tiem", "uong", "acid",
        }
        ingredient_tokens = sorted({t for t in tokens if len(t) >= 3 and t not in stop and not t.isdigit()})

    return {
        "doc_type": "drug_info",
        "drug_name": drug_name,
        "drug_name_upper": drug_name.upper(),
        "active_ingredient": active_ingredient,
        "active_ingredient_norm": _normalize_for_search(active_ingredient),
        "ingredient_keywords": ", ".join(ingredient_tokens),
        "price": price,
    }


def _split_consolidated_drug_catalog(doc: Document) -> List[Document]:
    """Split the large DANH_MUC_THUOC_NOI_BO_TOAN_BO.md into individual drug sections.

    Each section starts with '### NNN. DRUG_NAME' and ends before the next '### NNN.' or EOF.
    Part 1 (lookup table) and Part 2 (reverse lookup) are kept as general chunks.
    """
    content = doc.page_content
    source = doc.metadata.get("source", "DANH_MUC_THUOC_NOI_BO_TOAN_BO.md")
    base_meta = {k: v for k, v in doc.metadata.items()}

    # Split on "### NNN. " pattern (drug section headers in Part 3)
    pattern = re.compile(r"^### \d+\.\s+", re.MULTILINE)
    matches = list(pattern.finditer(content))

    sections: List[Document] = []

    # Part before first drug section (Parts 1 & 2 — lookup tables)
    if matches:
        preamble = content[:matches[0].start()].strip()
        if preamble:
            sections.append(Document(
                page_content=preamble,
                metadata={**base_meta, "section": "lookup_tables"},
            ))

    # Each drug section
    for i, m in enumerate(matches):
        start = m.start()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(content)
        section_text = content[start:end].strip()
        if section_text and len(section_text) > 50:
            sections.append(Document(
                page_content=section_text,
                metadata={**base_meta, "section": "drug_entry"},
            ))

    return sections


def _split_documents_with_drug_awareness(raw_docs: List[Document]) -> List[Document]:
    """Split docs while preserving key drug context in every drug chunk."""
    # Drug files: larger chunk + larger overlap to avoid losing dose/price/contraindication context.
    drug_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1800,
        chunk_overlap=350,
        separators=["\n## ", "\n**", "\n\n", "\n- ", ". ", "\n", " "],
    )
    # General docs keep current profile for speed.
    general_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)

    out: List[Document] = []
    for doc in raw_docs:
        source = str(doc.metadata.get("source", ""))

        # ── Special handling for consolidated drug catalog ──
        if "DANH_MUC_THUOC_NOI_BO" in source.upper():
            drug_sections = _split_consolidated_drug_catalog(doc)
            for section_doc in drug_sections:
                profile = _extract_drug_profile_from_content(section_doc.page_content)
                if profile:
                    section_doc.metadata.update(profile)
                    chunks = drug_splitter.split_documents([section_doc])
                    drug_name = profile.get("drug_name", "")
                    active = profile.get("active_ingredient", "")
                    keywords = profile.get("ingredient_keywords", "")
                    price = profile.get("price", "")
                    prefix_parts = [f"Thuoc: {drug_name}"]
                    if active:
                        prefix_parts.append(f"Hoat chat: {active}")
                    if price:
                        prefix_parts.append(f"Gia: {price}")
                    if keywords:
                        prefix_parts.append(f"Tu khoa hoat chat: {keywords}")
                    prefix = "[" + " | ".join(prefix_parts) + "]\n"
                    for ch in chunks:
                        if not ch.page_content.startswith("[Thuoc:"):
                            ch.page_content = prefix + ch.page_content
                        ch.metadata.update(profile)
                    out.extend(chunks)
                else:
                    out.extend(general_splitter.split_documents([section_doc]))
            logger.info("Consolidated drug catalog → %d drug sections", len(drug_sections))
            continue

        profile = _extract_drug_profile(doc)
        if profile:
            doc.metadata.update(profile)
            chunks = drug_splitter.split_documents([doc])

            drug_name = profile.get("drug_name", "")
            active = profile.get("active_ingredient", "")
            keywords = profile.get("ingredient_keywords", "")
            prefix_parts = [f"Thuoc: {drug_name}"]
            if active:
                prefix_parts.append(f"Hoat chat: {active}")
            if keywords:
                prefix_parts.append(f"Tu khoa hoat chat: {keywords}")
            prefix = "[" + " | ".join(prefix_parts) + "]\n"

            for ch in chunks:
                if not ch.page_content.startswith("[Thuoc:"):
                    ch.page_content = prefix + ch.page_content
                ch.metadata.update(profile)
            out.extend(chunks)
        else:
            out.extend(general_splitter.split_documents([doc]))

    return out


# ══════════════════════════════════════════════════════════════════════════════
#  HybridRAG Retriever Builder
# ══════════════════════════════════════════════════════════════════════════════

def _build_retrievers(data_path: str, db_path: str):
    """Build ChromaDB + BM25 Hybrid Retriever + CrossEncoder Reranker.

    Returns:
        (fast_retriever, deep_retriever, splits)
        fast_retriever: EnsembleRetriever (BM25 + Vector, k=15)
        deep_retriever: ContextualCompressionRetriever (Ensemble k=25 → Reranker top 5)
        splits: document splits for reference
    """
    from langchain_huggingface import HuggingFaceEmbeddings
    from langchain_chroma import Chroma
    from chromadb.config import Settings

    logger.info("--- Tải Embedding Model ---")
    embedding_model = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        model_kwargs={"device": DEVICE},
    )
    chroma_settings = Settings(anonymized_telemetry=False)

    # ChromaDB vector store — kiểm tra DB có sẵn TRƯỚC khi quét PDF
    import chromadb
    chroma_client = chromadb.PersistentClient(path=db_path, settings=chroma_settings)

    vectorstore = Chroma(
        client=chroma_client,
        embedding_function=embedding_model,
    )

    # Kiểm tra DB đã có dữ liệu chưa
    is_empty = True
    try:
        existing_count = len(vectorstore.get()["ids"])
        is_empty = existing_count == 0
        if not is_empty:
            logger.info("--- ChromaDB đã có %d documents sẵn tại %s ---", existing_count, db_path)
    except Exception:
        is_empty = True

    # Nếu đã có dữ liệu VÀ không cần rebuild → bỏ qua hoàn toàn việc quét PDF
    splits = []
    if not is_empty and not FORCE_REBUILD_DB:
        logger.info("--- Bỏ qua quét tài liệu, dùng ChromaDB có sẵn ---")
        # Lấy documents từ ChromaDB để khởi tạo BM25 (keyword search)
        try:
            from langchain_core.documents import Document
            chroma_data = vectorstore.get(include=["documents", "metadatas"])
            splits = [
                Document(page_content=doc, metadata=meta or {})
                for doc, meta in zip(chroma_data["documents"], chroma_data["metadatas"])
                if doc  # bỏ qua document rỗng
            ]
            logger.info("--- Đã lấy %d documents từ ChromaDB cho BM25 ---", len(splits))
        except Exception as e:
            logger.warning("Không thể lấy documents từ ChromaDB cho BM25: %s", e)
    else:
        # Chỉ quét PDF khi thực sự cần build lại
        raw_docs = _load_documents_from_folder(data_path)
        if raw_docs:
            splits = _split_documents_with_drug_awareness(raw_docs)
            logger.info("Split vào %d chunks.", len(splits))

        if not splits:
            if is_empty:
                logger.warning("⚠️ Không có dữ liệu để tạo Index!")
                return None, None, []
        else:
            logger.info("--- Tạo Index dữ liệu mới vào ChromaDB ---")
            if FORCE_REBUILD_DB and os.path.exists(db_path):
                shutil.rmtree(db_path, ignore_errors=True)
                chroma_client = chromadb.PersistentClient(path=db_path, settings=chroma_settings)

            vectorstore = Chroma(
                client=chroma_client,
                embedding_function=embedding_model,
            )

            # Thêm document theo từng khối (batch) để tránh lỗi batch size limit (5461)
            batch_size = 5000
            for i in range(0, len(splits), batch_size):
                batch = splits[i:i + batch_size]
                vectorstore.add_documents(batch)
                logger.info("Đã mã hóa và đưa %d documents vào ChromaDB (Tiến độ: %d/%d)", len(batch), min(i + batch_size, len(splits)), len(splits))

    # ── Fast Retriever: BM25 + Vector (k=15) ──────────────────────────────
    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 15})
    fast_retriever = vector_retriever

    if splits:
        try:
            from langchain_community.retrievers import BM25Retriever
            from langchain.retrievers.ensemble import EnsembleRetriever

            bm25_retriever = BM25Retriever.from_documents(splits, preprocess_func=_bm25_preprocess)
            bm25_retriever.k = 15
            fast_retriever = EnsembleRetriever(
                retrievers=[bm25_retriever, vector_retriever],
                # Slightly favor BM25 so ingredient keyword queries map to exact chunks.
                weights=[0.65, 0.35],
            )
            logger.info("Fast Retriever: BM25(%d docs) + Vector, k=15", len(splits))
        except Exception as e:
            logger.warning("BM25 init failed for fast, using vector-only: %s", e)

    # ── Deep Retriever: BM25 + Vector (k=25) → Reranker (top 5) ──────────
    vector_retriever_deep = vectorstore.as_retriever(search_kwargs={"k": 25})
    ensemble_deep = vector_retriever_deep

    if splits:
        try:
            from langchain_community.retrievers import BM25Retriever
            from langchain.retrievers.ensemble import EnsembleRetriever

            bm25_deep = BM25Retriever.from_documents(splits, preprocess_func=_bm25_preprocess)
            bm25_deep.k = 25
            ensemble_deep = EnsembleRetriever(
                retrievers=[bm25_deep, vector_retriever_deep],
                weights=[0.6, 0.4],
            )
        except Exception as e:
            logger.warning("BM25 init failed for deep: %s", e)

    # CrossEncoder Reranker
    deep_retriever = ensemble_deep
    try:
        from langchain_community.cross_encoders import HuggingFaceCrossEncoder
        from langchain.retrievers.document_compressors import CrossEncoderReranker
        from langchain.retrievers import ContextualCompressionRetriever

        logger.info("--- Tải Reranker Model (BGE-reranker-v2-m3) trên [%s] ---", DEVICE)
        reranker_model = HuggingFaceCrossEncoder(
            model_name="BAAI/bge-reranker-v2-m3",
            model_kwargs={"device": DEVICE},
        )
        compressor = CrossEncoderReranker(model=reranker_model, top_n=5)
        deep_retriever = ContextualCompressionRetriever(
            base_compressor=compressor,
            base_retriever=ensemble_deep,
        )
        logger.info("Deep Retriever: Ensemble(k=25) → Reranker(top_n=5)")
    except Exception as e:
        logger.warning("Reranker init failed, deep mode = ensemble only: %s", e)

    return fast_retriever, deep_retriever, splits


# ══════════════════════════════════════════════════════════════════════════════
#  Reference Helper
# ══════════════════════════════════════════════════════════════════════════════

def _clean_source_name(raw: str) -> str:
    """Convert raw filename to human-readable display name."""
    fname = raw.replace("\\", "/").split("/")[-1]
    display = fname
    # Bỏ hash dài (MD5-like) khỏi tên file markdown
    display = re.sub(r'\s*[0-9a-f]{20,}', '', display)
    # Bỏ đuôi file
    for ext in ('.md', '.pdf', '.docx', '.xlsx', '.txt', '.csv'):
        if display.lower().endswith(ext):
            display = display[:-len(ext)]
            break
    display = display.strip().strip('_').strip()
    return display if display else fname


def _extract_used_sources(answer: str) -> tuple:
    """Extract [USED_SOURCES: ...] tag from LLM answer.
    
    Returns:
        (clean_answer, used_source_names)
    """
    match = re.search(r'\[USED_SOURCES:\s*(.+?)\]', answer, re.IGNORECASE)
    if not match:
        return answer, []
    
    raw_sources = match.group(1).strip()
    # Xóa tag khỏi answer hiển thị
    clean_answer = answer[:match.start()].rstrip()
    
    if raw_sources.upper() == "NONE":
        return clean_answer, []
    
    # Parse danh sách tên file
    source_names = [s.strip() for s in raw_sources.split("|") if s.strip()]
    return clean_answer, source_names


def _build_references_text(docs: list, used_source_names: list = None) -> str:
    """Build clean Markdown reference block from retrieved documents.
    
    If used_source_names is provided, only include documents whose
    source filename matches one of the used sources (fuzzy match).
    """
    seen = set()
    sources = []
    
    for doc in docs:
        raw = doc.metadata.get("source", "")
        if not raw:
            continue
        fname = raw.replace("\\", "/").split("/")[-1]
        if fname in seen:
            continue
        
        display = _clean_source_name(raw)
        
        # Nếu có danh sách USED_SOURCES → chỉ lấy nguồn LLM thực sự dùng
        if used_source_names:
            is_used = False
            for used_name in used_source_names:
                used_clean = used_name.strip()
                # Fuzzy match: tên file chứa tên nguồn hoặc ngược lại
                if (used_clean.lower() in fname.lower() or 
                    used_clean.lower() in display.lower() or
                    fname.lower() in used_clean.lower() or
                    display.lower() in used_clean.lower()):
                    is_used = True
                    break
            if not is_used:
                continue
        
        seen.add(fname)
        if not display:
            display = fname
        sources.append(display)

    if not sources:
        return ""

    # Tối đa 5 nguồn tham khảo
    MAX_REFS = 5
    lines = [f"  - {s}" for s in sources[:MAX_REFS]]
    if len(sources) > MAX_REFS:
        lines.append(f"  - *...và {len(sources) - MAX_REFS} tài liệu khác*")
    return "\n".join(lines)


# ══════════════════════════════════════════════════════════════════════════════
#  DeepMedBot — Core Chat Engine
# ══════════════════════════════════════════════════════════════════════════════

class DeepMedBot:
    """HybridRAG chatbot: ChromaDB + BM25 + Reranker with 2 modes (Fast/Deep)."""

    def __init__(self):
        self.fast_chain = None
        self.deep_chain = None
        self.ready = False
        self.llm = None

        if not GOOGLE_API_KEYS and not NVIDIA_API_KEY and not OPENROUTER_API_KEY:
            logger.error("⚠️ Thiếu API key! Cần GOOGLE_API_KEY, GOOGLE_API_KEYS, NVIDIA_API_KEY hoặc OPENROUTER_API_KEY")
            return

        try:
            # Build retrievers
            self.fast_retriever, self.deep_retriever, self.splits = _build_retrievers(
                DATA_DIR, CHROMA_DB_PATH
            )

            # LLM with failover (Gemini → Qwen/OpenRouter)
            self.llm = get_llm()
            if not self.llm:
                logger.error("❌ Không thể khởi tạo LLM!")
                return
            logger.info("LLM active: %s", get_llm_name())

            if self.fast_retriever and self.deep_retriever:
                self._build_chains()
                self.ready = True
                logger.info("✅ DeepMed-AI sẵn sàng với 2 chế độ (Fast + Deep)!")
            else:
                logger.warning("⚠️ Không có dữ liệu. Bot sẽ chỉ dùng kiến thức nền.")
                self.ready = True

        except Exception as e:
            logger.error("🔥 Lỗi khởi tạo bot: %s", e)
            logger.debug(traceback.format_exc())

    def _build_chains(self):
        """Build LangChain retrieval chains for fast and deep modes."""

        # History-aware query rewrite prompt
        context_system_prompt = (
            "Dựa trên lịch sử chat và câu hỏi mới nhất, hãy viết lại câu hỏi "
            "thành một câu hoàn chỉnh mang tính chuyên môn y khoa để tìm kiếm thông tin. "
            "CHỈ TRẢ VỀ CÂU HỎI ĐÃ VIẾT LẠI, KHÔNG TRẢ LỜI."
        )
        context_prompt = ChatPromptTemplate.from_messages([
            ("system", context_system_prompt),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ])

        # Document citation template
        document_prompt = PromptTemplate(
            input_variables=["page_content", "source"],
            template="[Trích từ tài liệu: {source}]\n{page_content}",
        )

        # Fast mode QA prompt
        qa_prompt_fast = ChatPromptTemplate.from_messages([
            ("system", QA_SYSTEM_PROMPT_FAST),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ])

        # Deep mode QA prompt (with 3-step clinical reasoning)
        qa_prompt_deep = ChatPromptTemplate.from_messages([
            ("system", QA_SYSTEM_PROMPT_DEEP),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ])

        # Build QA chains (shared between history-aware and direct paths)
        self._qa_chain_fast = create_stuff_documents_chain(
            self.llm, qa_prompt_fast, document_prompt=document_prompt,
        )
        self._qa_chain_deep = create_stuff_documents_chain(
            self.llm, qa_prompt_deep, document_prompt=document_prompt,
        )

        # Direct chains (no history rewrite — saves 1 LLM call)
        self.fast_chain_direct = create_retrieval_chain(self.fast_retriever, self._qa_chain_fast)
        self.deep_chain_direct = create_retrieval_chain(self.deep_retriever, self._qa_chain_deep)

        # History-aware chains (with query rewrite)
        history_aware_fast = create_history_aware_retriever(
            self.llm, self.fast_retriever, context_prompt,
        )
        self.fast_chain = create_retrieval_chain(history_aware_fast, self._qa_chain_fast)

        history_aware_deep = create_history_aware_retriever(
            self.llm, self.deep_retriever, context_prompt,
        )
        self.deep_chain = create_retrieval_chain(history_aware_deep, self._qa_chain_deep)

    def chat(self, message: str, chat_history: list, mode: str = "fast") -> Dict[str, Any]:
        """Process a message and return response + source info.

        Args:
            message: User's question
            chat_history: List of {"role": ..., "content": ...}
            mode: "fast" or "deep"

        Returns:
            {"response": str, "source": str}
        """
        if not self.ready:
            return {
                "response": "⏳ Hệ thống đang khởi động hoặc gặp lỗi cấu hình...",
                "source": "Thông báo hệ thống",
            }

        # Convert history to LangChain messages
        lc_history = []
        if chat_history:
            for item in chat_history[-MAX_HISTORY_TURNS * 2:]:
                role = item.get("role", "")
                content = item.get("content", "")
                if role == "user" and content.strip():
                    lc_history.append(HumanMessage(content=content))
                elif role == "assistant" and content.strip():
                    lc_history.append(AIMessage(content=content))

        # Select chain — skip history rewrite when no conversation history
        has_history = bool(lc_history)
        if mode == "deep":
            active_chain = self.deep_chain if has_history else self.deep_chain_direct
        else:
            active_chain = self.fast_chain if has_history else self.fast_chain_direct
        mode_label = "Chuyên sâu (Reranker)" if mode == "deep" else "Tốc độ"
        if not has_history:
            logger.info("No history → using direct chain (skip query rewrite)")

        if not active_chain:
            # No data → fallback to LLM only
            try:
                resp = self.llm.invoke([HumanMessage(content=message)])
                return {
                    "response": f"⚠️ *(Chế độ kiến thức mở — Không có dữ liệu nội bộ)*\n\n{resp.content}",
                    "source": f"Kiến thức y khoa tổng quát ({get_llm_name()})",
                }
            except Exception:
                return {
                    "response": "❌ Lỗi: Không thể kết nối với AI. Vui lòng kiểm tra API Key.",
                    "source": "Thông báo hệ thống",
                }

        # Run RAG chain
        try:
            result = active_chain.invoke({"input": message, "chat_history": lc_history})
            raw_answer = result.get("answer", "")
            retrieved_docs = result.get("context", [])

            # Tách tag [USED_SOURCES: ...] khỏi câu trả lời
            answer, used_source_names = _extract_used_sources(raw_answer)
            logger.info("LLM cited sources: %s", used_source_names)

            # Build source citation — CHỈ từ nguồn LLM thực sự dùng
            if retrieved_docs:
                refs = _build_references_text(retrieved_docs, used_source_names)
                source_text = f"Dữ liệu nội bộ ({mode_label})"
                if refs:
                    answer += f"\n\n---\n\n📚 **Nguồn tham khảo:**\n{refs}"
            else:
                source_text = "Cơ sở dữ liệu y tế nội bộ"

            logger.info("Chat OK: mode=%s, docs=%d, used_sources=%d, answer_len=%d",
                        mode, len(retrieved_docs), len(used_source_names), len(answer))
            return {"response": answer, "source": source_text}

        except Exception as e:
            logger.error("Chat error: %s", e)
            logger.error(traceback.format_exc())

            # Fallback to LLM only
            try:
                resp = self.llm.invoke([HumanMessage(content=message)])
                return {
                    "response": f"⚠️ Quá trình trích xuất RAG gặp lỗi. "
                                f"Đang dùng kiến thức nền:\n\n{resp.content}",
                    "source": "Kiến thức y khoa tổng quát (Fallback)",
                }
            except Exception:
                return {
                    "response": f"❌ Đã xảy ra lỗi hệ thống. Vui lòng thử lại. (Mã lỗi: {str(e)})",
                    "source": "Thông báo hệ thống",
                }

    async def chat_stream(self, message: str, chat_history: list, mode: str = "fast"):
        """Stream response tokens via async generator.

        Yields dicts: {"type": "token"|"done", ...}
        Retrieval is done upfront (sync); only LLM generation is streamed.
        """
        if not self.ready:
            yield {"type": "token", "content": "⏳ Hệ thống đang khởi động hoặc gặp lỗi cấu hình..."}
            yield {"type": "done", "source": "Thông báo hệ thống", "refs": ""}
            return

        # Convert history
        lc_history = []
        if chat_history:
            for item in chat_history[-MAX_HISTORY_TURNS * 2:]:
                role = item.get("role", "")
                content = item.get("content", "")
                if role == "user" and content.strip():
                    lc_history.append(HumanMessage(content=content))
                elif role == "assistant" and content.strip():
                    lc_history.append(AIMessage(content=content))

        has_history = bool(lc_history)
        mode_label = "Chuyên sâu (Reranker)" if mode == "deep" else "Tốc độ"

        # Pick retriever and QA chain
        if mode == "deep":
            retriever = self.deep_retriever
            qa_chain = self._qa_chain_deep
        else:
            retriever = self.fast_retriever
            qa_chain = self._qa_chain_fast

        if not retriever or not qa_chain:
            yield {"type": "token", "content": "⚠️ Không có dữ liệu nội bộ."}
            yield {"type": "done", "source": "Thông báo hệ thống", "refs": ""}
            return

        try:
            # Step 1: Rewrite query if history exists (sync — fast)
            search_query = message
            if has_history:
                try:
                    rewrite_prompt = (
                        "Dựa trên lịch sử chat, viết lại câu hỏi thành câu hoàn chỉnh để tìm kiếm. "
                        "CHỈ TRẢ VỀ CÂU HỎI, KHÔNG TRẢ LỜI.\n\n"
                        f"Lịch sử: {[m.content[:100] for m in lc_history[-4:]]}\n"
                        f"Câu hỏi mới: {message}"
                    )
                    rewrite_resp = self.llm.invoke([HumanMessage(content=rewrite_prompt)])
                    search_query = rewrite_resp.content.strip() or message
                    logger.info("Rewritten query: %s", search_query[:80])
                except Exception:
                    search_query = message

            # Step 2: Retrieve documents (sync — fast)
            retrieved_docs = retriever.invoke(search_query)
            logger.info("Retrieved %d docs for streaming", len(retrieved_docs))

            # Step 3: Stream LLM generation
            full_answer = []
            async for chunk in qa_chain.astream({
                "input": message,
                "chat_history": lc_history,
                "context": retrieved_docs,
            }):
                token = chunk if isinstance(chunk, str) else str(chunk)
                if token:
                    full_answer.append(token)
                    yield {"type": "token", "content": token}

            # Step 4: Post-process (refs)
            raw_answer = "".join(full_answer)
            answer, used_source_names = _extract_used_sources(raw_answer)
            refs = ""
            source_text = f"Dữ liệu nội bộ ({mode_label})"
            if retrieved_docs:
                refs = _build_references_text(retrieved_docs, used_source_names)

            yield {"type": "done", "source": source_text, "refs": refs, "answer": answer}

        except Exception as e:
            logger.error("Stream error: %s", e)
            logger.error(traceback.format_exc())
            yield {"type": "token", "content": f"❌ Đã xảy ra lỗi: {str(e)}"}
            yield {"type": "done", "source": "Thông báo hệ thống", "refs": ""}

class ChatService:
    """Orchestrates DeepMedBot for FastAPI endpoints."""

    def __init__(self):
        self.bot: Optional[DeepMedBot] = None
        self.conversation_histories: Dict[str, list] = {}
        self.workflow_app = True  # Compatibility flag for chat endpoint check
        logger.info("ChatService initialized")

    def initialize(self) -> None:
        """Initialize the DeepMedBot (called once at startup)."""
        logger.info("Initializing DeepMedBot...")
        self.bot = DeepMedBot()
        if self.bot.ready:
            self.workflow_app = True
            logger.info("DeepMedBot initialized successfully")
        else:
            self.workflow_app = None
            logger.error("DeepMedBot initialization failed")

    async def process_message(self, session_id: str, message: str) -> Dict[str, Any]:
        """Run the HybridRAG pipeline for a single user message."""
        logger.info("Processing message for session %s...", session_id[:8])

        if not self.bot or not self.bot.ready:
            raise ValueError("Bot not initialized")

        # Persist user message
        db_service.save_message(session_id, "user", message)

        # Get conversation history for this session
        if session_id not in self.conversation_histories:
            self.conversation_histories[session_id] = []

        history = self.conversation_histories[session_id]

        # Auto-detect mode
        mode = self._detect_mode(message)

        # Run chat
        result = self.bot.chat(message, history, mode=mode)

        response_text = result["response"]
        source = result["source"]

        # Update in-memory history
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response_text})
        # Keep only last N turns
        if len(history) > MAX_HISTORY_TURNS * 2:
            self.conversation_histories[session_id] = history[-MAX_HISTORY_TURNS * 2:]

        # Persist assistant response
        db_service.save_message(session_id, "assistant", response_text, source)

        return {
            "response": response_text,
            "source": source,
            "timestamp": datetime.now().strftime("%I:%M %p"),
            "success": bool(response_text),
        }

    def clear_conversation(self, session_id: str) -> None:
        """Reset the in-memory conversation history for a session."""
        if session_id in self.conversation_histories:
            self.conversation_histories[session_id] = []
            logger.info("Conversation cleared for session %s", session_id[:8])

    def _detect_mode(self, message: str) -> str:
        """Auto-detect fast/deep mode based on keywords."""
        q_lower = message.lower()
        deep_indicators = [
            "phác đồ", "điều trị", "chẩn đoán", "xử trí", "hướng dẫn",
            "protocol", "guideline", "bệnh nhân bị", "cách điều trị",
            "dùng thuốc gì", "nên dùng", "kê đơn", "toa thuốc",
            "sơ cứu", "cấp cứu",
            "cảnh giác dược", "tương tác thuốc", "adr", "tác dụng phụ",
            "chống chỉ định", "hiệu chỉnh liều", "suy thận", "suy gan",
            "thai kỳ", "cho con bú", "người cao tuổi",
            "kháng sinh", "phối hợp thuốc", "thay thế thuốc",
            "khuyến cáo", "bộ y tế", "hiệp hội", "who",
        ]
        if any(ind in q_lower for ind in deep_indicators):
            logger.info("Auto-detected Deep mode for: %s", message[:50])
            return "deep"
        return "fast"

    async def process_message_stream(self, session_id: str, message: str):
        """Stream HybridRAG response as SSE events.

        Yields SSE-formatted strings: 'data: {...}\n\n'
        """
        import json as _json

        logger.info("Stream processing for session %s...", session_id[:8])

        if not self.bot or not self.bot.ready:
            yield f"data: {_json.dumps({'type': 'token', 'content': '⏳ Hệ thống đang khởi động...'})}\n\n"
            yield f"data: {_json.dumps({'type': 'done', 'source': 'Thông báo hệ thống'})}\n\n"
            return

        db_service.save_message(session_id, "user", message)

        if session_id not in self.conversation_histories:
            self.conversation_histories[session_id] = []
        history = self.conversation_histories[session_id]

        mode = self._detect_mode(message)

        full_answer = ""
        source = ""
        async for event in self.bot.chat_stream(message, history, mode=mode):
            if event["type"] == "token":
                yield f"data: {_json.dumps(event, ensure_ascii=False)}\n\n"
            elif event["type"] == "done":
                source = event.get("source", "")
                refs = event.get("refs", "")
                full_answer = event.get("answer", "")
                # Send cleaned answer (without [USED_SOURCES:...] tag) so frontend can replace
                clean_content = full_answer
                if refs:
                    clean_content += f"\n\n---\n\n📚 **Nguồn tham khảo:**\n{refs}"
                yield f"data: {_json.dumps({'type': 'replace', 'content': clean_content}, ensure_ascii=False)}\n\n"
                yield f"data: {_json.dumps({'type': 'done', 'source': source}, ensure_ascii=False)}\n\n"

        # Update history
        if full_answer:
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": full_answer})
            if len(history) > MAX_HISTORY_TURNS * 2:
                self.conversation_histories[session_id] = history[-MAX_HISTORY_TURNS * 2:]
            db_service.save_message(session_id, "assistant", full_answer, source)


# Module-level singleton
chat_service = ChatService()