backend/app/agents/executor.py

"""
DeepMed-AI — agents/executor.py
ExecutorAgent: synthesizes final response from RAG documents or LLM knowledge.
Vietnamese medical system prompt — TTYT Khu vực Thanh Ba.
"""

import re

from app.core.logging_config import logger
from app.core.state import AgentState
from app.tools.llm_client import DEEPMED_SYSTEM_PROMPT, get_llm


def _clean_drug_filename(fname: str) -> str:
    """Extract clean drug name from .md filename.

    Examples:
        "MIDANTIN 2010bc1073c281da9cfad7d9ce373065.md" → "MIDANTIN"
        "A T Esomeprazol 20 inj.md" → "A T Esomeprazol 20 inj"
        "Ceftriaxone 2000 2010bc1073c281fca190c0ea0ba09d46.md" → "Ceftriaxone 2000"
    """
    # Remove .md extension
    name = fname.replace(".md", "").strip()
    # Remove Notion-style hex hash (20+ hex chars at end, preceded by space)
    name = re.sub(r'\s+[0-9a-f]{20,}$', '', name, flags=re.IGNORECASE)
    return name.strip()


def _build_citations(documents: list) -> str:
    """Build a clean, grouped citation string from retrieved documents.
    
    Groups by source type:
    - Drug docs (.md): Just drug name
    - PDFs: Filename + page range (not per-page)
    """
    drug_cites = []
    pdf_cites = {}  # filename -> set of pages
    other_cites = []
    
    for doc in documents:
        fp = doc.metadata.get("source", "")
        fname = fp.split("/")[-1].split("\\")[-1]
        
        if fname.endswith(".md"):
            clean = _clean_drug_filename(fname)
            if clean and clean not in drug_cites:
                drug_cites.append(clean)
        elif fname.endswith(".pdf"):
            page = doc.metadata.get("page", "")
            if fname not in pdf_cites:
                pdf_cites[fname] = set()
            if page != "":
                pdf_cites[fname].add(str(page))
        elif fname:
            if fname not in other_cites:
                other_cites.append(fname)
    
    # Build clean output
    parts = []
    if drug_cites:
        parts.append(", ".join(drug_cites) + " (Thông tin thuốc nội bộ)")
    for pdf_name, pages in pdf_cites.items():
        if pages:
            sorted_pages = sorted(pages, key=lambda x: int(x) if x.isdigit() else 0)
            parts.append(f"{pdf_name} (tr.{','.join(sorted_pages)})")
        else:
            parts.append(pdf_name)
    for c in other_cites:
        parts.append(c)
    
    return "📚 " + " | ".join(parts) if parts else "Cơ sở dữ liệu y tế nội bộ"


def _get_response_instructions(query_type: str) -> str:
    """Return query-type-specific LLM instructions.
    
    Inspired by the reference chatbot's dual-mode approach:
    - Fast mode: Direct answer extraction from context
    - Deep mode: 3-step clinical reasoning (scan protocol -> cross-ref drugs -> synthesize)
    """
    
    if query_type == "ingredient_list":
        return (
            "**HUONG DAN TRA LOI (LIET KE THUOC THEO HOAT CHAT):**\n\n"
            "Nguoi dung dang hoi ve MOT HOAT CHAT. Hay tim TAT CA thuoc trong Context chua hoat chat do.\n\n"
            "1. Quet tat ca tai lieu trong Context. Tim cac header \"[Thuoc: ...]\" hoac \"[Trich tu tai lieu: ...]\".\n"
            "2. LIET KE TAT CA thuoc chua hoat chat duoc hoi duoi dang bang:\n"
            "   | STT | Ten thuoc | Hoat chat & Ham luong | Gia (VND) | Hang san xuat |\n"
            "3. Sau bang, tom gon su khac biet giua cac thuoc (dang bao che, duong dung...).\n"
            "4. Neu chi tim thay 1 thuoc -> liet ke va ghi ro \"Trong danh muc noi bo hien co 1 thuoc...\".\n"
            "5. KHONG bia them thuoc ngoai Context."
        )

    elif query_type == "protocol":
        return (
            "**HUONG DAN TRA LOI (PHAC DO DIEU TRI / CHAN DOAN):**\n\n"
            "QUY TRINH TU DUY (suy luan noi bo, KHONG in ra):\n"
            "- Buoc 1: Quet Context tim Huong dan dieu tri / Phac do lien quan.\n"
            "- Buoc 2: Doi chieu xem trong Context co san thuoc/hoat chat nao dap ung phac do.\n"
            "- Buoc 3: Duc ket thanh loi khuyen toan dien dua tren CHINH XAC nhung thuoc co san.\n\n"
            "TRINH BAY KET QUA theo 2 PHAN:\n\n"
            "**PHAN 1 - PHAC DO CHUAN:**\n"
            "- Phac do theo huong dan BYT Viet Nam / WHO.\n"
            "- Neu ro: cac bac dieu tri, thuoc lua chon (hoat chat, lieu, duong dung).\n"
            "- Truong hop dac biet: tre em, thai phu, nguoi gia, suy than/gan.\n\n"
            "**PHAN 2 - TOI UU VOI DANH MUC THUOC TTYT THANH BA:**\n"
            "BAT BUOC: Phai doc tat ca tai lieu trong Context de tim thuoc co san.\n"
            "Voi moi buoc trong phac do, tim CHINH XAC thuoc nao trong Context tuong ung:\n"
            "  - Tim header [Trich tu tai lieu: TEN THUOC (Danh muc thuoc)] trong Context\n"
            "  - Trich xuat: ten thuoc, ham luong, gia, duong dung tu Context\n"
            "  - Trinh bay BANG CU THE:\n"
            "    | Buoc phac do | Hoat chat can | Thuoc co san (tu Context) | Ham luong | Gia (VND) | Lieu dung cu the |\n"
            "  - Neu Context KHONG co thuoc tuong ung -> ghi ro \"Chua co trong danh muc\"\n"
            "  - Neu Context co nhieu thuoc cung hoat chat -> liet ke tat ca va so sanh\n\n"
            "TUYET DOI KHONG tra loi chung chung kieu 'tham khao danh muc'. "
            "PHAI chi ra DICH DANH ten thuoc, gia, lieu luong cu the tu Context.\n"
            "Viet tu nhien, mach lac, giong dong nghiep trao doi chuyen mon."
        )

    elif query_type == "drug_info":
        return (
            "**HUONG DAN TRA LOI (THONG TIN THUOC CU THE):**\n"
            "1. Xac dinh thuoc dung qua header \"[Thuoc: ...]\" hoac \"[Trich tu tai lieu: ...]\".\n"
            "2. Trich xuat va tra loi TAT CA thong tin co lien quan:\n"
            "   - Ten thuoc, hoat chat, ham luong, dang bao che\n"
            "   - **Gia** (neu Context co dong \"Gia cua thuoc...\" -> BAT BUOC tra loi gia)\n"
            "   - Hang san xuat, so dang ky\n"
            "   - Chi dinh, chong chi dinh, lieu dung, tac dung phu, tuong tac thuoc\n"
            "3. Neu Context co tai lieu ve thuoc KHAC (hoi thuoc A nhung Context co thuoc B) -> BO QUA.\n"
            "4. TUYET DOI KHONG lay thong tin thuoc B gan cho thuoc A."
        )

    else:  # general
        return (
            "**HUONG DAN TRA LOI:**\n"
            "1. Doc Context, trich xuat thong tin lien quan den cau hoi.\n"
            "2. Neu Context khong du -> dung kien thuc y khoa chung, ghi ro nguon.\n"
            "3. Tra loi tu nhien, mach lac, giong dong nghiep trao doi chuyen mon.\n"
            "4. Khuyen nghi tham kham truc tiep khi can."
        )


def ExecutorAgent(state: AgentState) -> AgentState:
    """Synthesize final response from RAG documents or LLM knowledge."""
    llm = get_llm()
    question = state["question"]
    source_info = state.get("source", "Unknown")

    # ── Conversation history (last 5 turns) ────────────────────────────────
    history_lines = []
    for item in state.get("conversation_history", [])[-5:]:
        role = "Bác sĩ/Người dùng" if item.get("role") == "user" else "DeepMed-AI"
        history_lines.append(f"{role}: {item.get('content', '')}")
    history_context = "\n".join(history_lines) or "Không có lịch sử."

    # ── Case 0: LLM unavailable ────────────────────────────────────────────
    if not llm:
        answer = (
            "⚠️ Hệ thống AI tạm thời không khả dụng. "
            "Vui lòng liên hệ trực tiếp bác sĩ hoặc gọi 115."
        )
        source_info = "Thông báo hệ thống"

    # ── Case 1: Have documents from Vector Store ──────────────────
    elif state.get("documents") and len(state["documents"]) > 0:
        # Format documents with citation tags matching reference chatbot style
        max_docs = 8 if state.get("query_type") == "protocol" else 6
        sections = []
        for i, doc in enumerate(state["documents"][:max_docs], 1):
            src = doc.metadata.get("source", "Không rõ")
            src_short = src.split("/")[-1].split("\\")[-1]
            drug = doc.metadata.get("drug_name_original", "")
            
            # Build citation tag like reference chatbot's document_prompt
            if drug:
                cite_tag = f"{drug} (Thông tin thuốc nội bộ)"
            elif src_short.endswith(".md"):
                cite_tag = f"{_clean_drug_filename(src_short)} (Danh mục thuốc)"
            else:
                page = doc.metadata.get("page", "")
                cite_tag = f"{src_short} (tr.{page})" if page else src_short
            
            sections.append(f"[Trích từ tài liệu: {cite_tag}]\n{doc.page_content[:3000]}")

        content = "\n\n---\n\n".join(sections)
        source_info = _build_citations(state["documents"])

        prompt = f"""{DEEPMED_SYSTEM_PROMPT}

---

## Lịch sử hội thoại:
{history_context}

## Câu hỏi hiện tại:
{question}

## Tài liệu y tế từ CSDL nội bộ (tham khảo):
{content}

---

{_get_response_instructions(state.get("query_type", "general"))}"""

        try:
            response = llm.invoke(prompt)
            answer = (
                response.content.strip()
                if hasattr(response, "content")
                else str(response).strip()
            )
            logger.info("Executor: Generated from %d documents", len(state["documents"]))
        except Exception as e:
            logger.error("Executor: LLM failed: %s", e)
            answer = (
                "Tôi hiểu câu hỏi của bạn. Để được tư vấn chính xác, "
                "vui lòng tham khảo ý kiến bác sĩ tại TTYT Thanh Ba."
            )
            source_info = "Thông báo hệ thống"

    # ── Case 2: Pre-generated LLM response (from LLMAgent) ────────────────
    elif state.get("llm_success") and state.get("generation"):
        answer = state["generation"]
        logger.info("Executor: Using pre-generated LLM response")

    # ── Case 3: Final fallback ─────────────────────────────────────────────
    else:
        answer = (
            "Tôi chưa tìm được thông tin phù hợp trong cơ sở dữ liệu nội bộ. "
            "Vui lòng mô tả chi tiết hơn hoặc liên hệ bác sĩ tại TTYT Khu vực Thanh Ba."
        )
        source_info = "Thông báo hệ thống"

    state["generation"] = answer
    state["source"] = source_info
    state["conversation_history"].append({"role": "user", "content": question})
    state["conversation_history"].append(
        {"role": "assistant", "content": answer, "source": source_info}
    )
    return state