backend/app/tools/document_loader.py

"""
BaoThang-AI — tools/document_loader.py
Multi-format document loading with drug-aware chunking.
Every chunk from a drug .md file:
  1. Is prefixed with [Thuốc: NAME | Hoạt chất: ...] for embedding search
  2. Has metadata drug_name="NAME" + doc_type="drug_info" for FILTERED search
This ensures 100% retrieval accuracy for any drug query.
"""

import os
import re
import glob
from typing import List, Optional, Set, Tuple
from langchain_core.documents import Document
from app.core.logging_config import logger


# ── File Loaders ───────────────────────────────────────────────────────────────

def load_pdf(file_path: str) -> List[Document]:
    from langchain_community.document_loaders import PyPDFLoader
    return PyPDFLoader(file_path).load()


def load_docx(file_path: str) -> List[Document]:
    import docx
    try:
        doc = docx.Document(file_path)
        content = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
        if content:
            return [Document(page_content=content, metadata={"source": file_path})]
    except Exception as e:
        logger.error("Failed to load DOCX %s: %s", file_path, e)
    return []


def load_text(file_path: str) -> List[Document]:
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
        if content.strip():
            return [Document(page_content=content, metadata={"source": file_path})]
    except Exception as e:
        logger.error("Failed to load Text/MD %s: %s", file_path, e)
    return []


def load_smart_excel(file_path: str) -> List[Document]:
    """Each Excel row becomes a separate Document."""
    import pandas as pd
    docs = []
    try:
        if file_path.lower().endswith(".csv"):
            df = pd.read_csv(file_path)
        else:
            df = pd.read_excel(file_path)
        for index, row in df.iterrows():
            items = [f"[{col}: {val}]" for col, val in row.items()
                     if pd.notna(val) and str(val).strip()]
            if items:
                docs.append(Document(
                    page_content=" | ".join(items),
                    metadata={"source": file_path, "row": index + 1},
                ))
    except Exception as e:
        logger.error("Failed to load Excel %s: %s", file_path, e)
    return docs


# ── Drug Info Helpers ──────────────────────────────────────────────────────────

def _extract_drug_header(content: str) -> Tuple[str, str, str]:
    """Extract drug name + active ingredient from the first lines of a .md file.

    Expected format:
        # MIDANTIN
        (blank line)
        Hoạt chất: Amoxicilin+acid clavulanic 1g+0,2g

    Returns:
        (header_prefix, drug_name, active_ingredient)
        header_prefix: "[Thuốc: MIDANTIN | Hoạt chất: Amoxicilin+acid clavulanic 1g+0,2g]\n"
        drug_name: "MIDANTIN"
        active_ingredient: "Amoxicilin+acid clavulanic 1g+0,2g"
    """
    lines = content.split("\n", 12)[:12]
    drug_name = ""
    active = ""

    for line in lines:
        s = line.strip()
        if s.startswith("# ") and not drug_name:
            drug_name = s[2:].strip()
        if not active and "hoạt chất" in s.lower():
            m = re.search(r'[Hh]oạt chất[:\s]+(.+)', s)
            if m:
                active = m.group(1).strip()

    if drug_name:
        parts = [f"Thuốc: {drug_name}"]
        if active:
            parts.append(f"Hoạt chất: {active}")
        prefix = "[" + " | ".join(parts) + "]\n"
        return prefix, drug_name, active
    return "", "", ""


def _parse_ingredient_keywords(active_ingredient: str) -> List[str]:
    """Parse active ingredient string into searchable uppercase keywords.

    Examples:
        "Amoxicilin+acid clavulanic 1g+0,2g"
        → ["AMOXICILIN", "CLAVULANIC"]

        "Ceftriaxon dưới dạng Ceftriaxon natri 2000mg"
        → ["CEFTRIAXON", "CEFTRIAXON", "NATRI"]

        "Paracetamol 500mg"
        → ["PARACETAMOL"]
    """
    if not active_ingredient:
        return []

    # Split on common separators: +, comma, semicolon, slash, space
    tokens = re.split(r'[+,;/()\s]+', active_ingredient)

    # Filter: keep words ≥3 chars, alphabetic, not dosage/unit numbers
    SKIP_WORDS = {
        "MG", "ML", "MCG", "IU", "DẠ", "DẠNG", "DƯỚI", "MỖI",
        "ACID", "VIÊN", "NÉN", "GÓI", "LỌ", "ỐNG", "LIỀU",
        "THUỐC", "TIÊM", "UỐNG", "HOẠT", "CHẤT",
    }
    keywords = []
    for t in tokens:
        t_clean = t.strip().upper()
        if len(t_clean) < 3:
            continue
        # Skip pure numbers or dosage patterns
        if re.match(r'^[\d.,]+$', t_clean):
            continue
        # Skip dosage with units like "1G", "200MG"
        if re.match(r'^\d+[A-Z]{1,3}$', t_clean):
            continue
        if t_clean in SKIP_WORDS:
            continue
        keywords.append(t_clean)

    return list(set(keywords))  # Deduplicate


def _is_drug_info_file(file_path: str) -> bool:
    """Detect drug info .md files by directory name.

    Works on both Windows (backslash) and Linux/Docker (forward slash).
    Handles both Unicode and ASCII-normalized folder names.
    """
    normed = file_path.replace("\\", "/").lower()
    # Primary check: Vietnamese Unicode folder name
    if "thông tin thuốc nội bộ" in normed:
        return True
    # Fallback: ASCII-normalized (in case Docker normalizes differently)
    if "thong tin thuoc noi bo" in normed:
        return True
    # Fallback: check by path segment containing "thuoc"
    parts = normed.split("/")
    for part in parts:
        if "thuoc" in part and "noi" in part and "bo" in part:
            return True
    return False


# ── Chunking ───────────────────────────────────────────────────────────────────

def split_documents(
    docs: List[Document],
    drug_prefix: str = "",
    extra_metadata: Optional[dict] = None,
) -> List[Document]:
    """Split docs into overlapping chunks.
    
    If drug_prefix is set, prepend it to every chunk.
    If extra_metadata is set, merge it into every chunk's metadata.
    """
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1024,
        chunk_overlap=128,
        separators=["\n\n", ". ", "\n", " "],
    )
    chunks = splitter.split_documents(docs)

    if drug_prefix:
        for chunk in chunks:
            # Only prepend if not already there (avoid double-prefix)
            if not chunk.page_content.startswith(drug_prefix):
                chunk.page_content = drug_prefix + chunk.page_content

    if extra_metadata:
        for chunk in chunks:
            chunk.metadata.update(extra_metadata)

    return chunks


# ── Main Loader ────────────────────────────────────────────────────────────────

def load_all_documents(directory: str) -> List[Document]:
    """Recursively load all supported files from the data directory.
    
    Drug .md files get special metadata: drug_name, active_ingredient, doc_type.
    This enables filtered search in Qdrant for 100% retrieval accuracy.
    """
    all_docs: List[Document] = []
    drug_count = 0

    loaders = {
        ".pdf": load_pdf,
        ".docx": load_docx,
        ".txt": load_text,
        ".csv": load_smart_excel,
        ".xlsx": load_smart_excel,
    }

    pattern = os.path.join(directory, "**", "*.*")
    files = glob.glob(pattern, recursive=True)
    logger.info("Scanning %s: found %d files", directory, len(files))

    for fp in files:
        ext = os.path.splitext(fp)[1].lower()

        # ── Skip individual drug .md files (thông tin thuốc nội bộ) ─────
        # All drug info is consolidated in DANH_MUC_THUOC_NOI_BO_TOAN_BO.md
        if ext == ".md" and _is_drug_info_file(fp):
            drug_count += 1
            continue

        # ── Non-drug .md files ──────────────────────────────────────────────
        if ext == ".md":
            docs = load_text(fp)
            if docs:
                meta = {"doc_type": "reference_md"}
                all_docs.extend(split_documents(docs, extra_metadata=meta))
            continue

        # ── PDF, DOCX, TXT, XLSX ────────────────────────────────────────────
        if ext in loaders:
            docs = loaders[ext](fp)
            if docs:
                meta = {"doc_type": f"reference_{ext.lstrip('.')}"}
                if ext != ".xlsx":
                    docs = split_documents(docs, extra_metadata=meta)
                else:
                    for d in docs:
                        d.metadata.update(meta)
                all_docs.extend(docs)

    logger.info(
        "Loaded %d chunks total (%d drug info files, %d other files)",
        len(all_docs), drug_count, len(files) - drug_count,
    )
    return all_docs