#!/usr/bin/env python3
"""
PMDA drug information MCP server — 真实 PG / OS 查询版本（替换原 mock）.

Plugin 自包含，不依赖 mygpt.* 任何模块。配置通过环境变量：
  PMDA_PG_HOST / PMDA_PG_PORT / PMDA_PG_DB / PMDA_PG_USER / PMDA_PG_PASSWORD
  PMDA_OPENSEARCH_URL (or OPENSEARCH_URL) / PMDA_OS_INDEX

参考 hu-sandbox/pmda/agent/tools.py 的 10 个 tool 行为（98/100 v2e baseline）.
"""

import asyncio
import sys
from dataclasses import asdict
from decimal import Decimal
from typing import Any, Dict, List, Optional, Sequence, Tuple

from mcp_common import (
    create_error_response,
    create_initialize_response,
    create_ping_response,
    create_tools_list_response,
    load_tools_from_json,
    handle_mcp_streaming,
)

from db import session
from queries import (
    drug_dosing_query,
    drug_interaction_query,
    drug_master_get,
    drug_restriction_query,
    list_categories_with_counts,
    list_drugs_in_category as _sql_list_drugs_in_category,
    search_drugs_in_db,
)
from os_client import client as os_client, INDEX_NAME as OS_INDEX_NAME


# ---------------------------------------------------------------------------
# Plain-text rendering (agent-friendly tool output)
#
# 工具结果以纯文本返回（而非 JSON），降低 agent 的解析负担与 token 噪音。
# CITATION enforcement 仍是工程化保证（不依赖 LLM 自觉）：
#   1. `_CITE_INSTRUCTION_TEXT` 注入每个含可引用源的结果头部 (LLM 第一眼)
#   2. 每条记录末尾一行 `CITATION:` 镜像 `_cite._tag` (LLM 直接复制, 不用 traverse)
#   3. `read_drug_chapter` 三明治包装 raw markdown (tag 物理紧贴章节文本)
# 命中 0 条时返回英文 no-results 话术, 且 **不含** CITATION 指令 —— 避免诱导
# agent 在无来源时编造引用。
# ---------------------------------------------------------------------------

_CITE_INSTRUCTION_TEXT = (
    "=== CITATION INSTRUCTIONS ===\n"
    "Each record below ends with a `CITATION:` line — a pre-built "
    "`<CITATION file=\"...\" filename=\"...\" />` tag the frontend PDF-highlight "
    "pipeline depends on. When you use a fact from a record, copy that record's "
    "`CITATION:` tag VERBATIM (byte-for-byte) immediately AFTER the paragraph or "
    "bullet that states the fact. NEVER collect citations at the end. At most ONE "
    "tag per unique file. Do NOT add, modify, reorder, remove attributes, or build "
    "a tag yourself. Records without a `CITATION:` line carry no clickable source — "
    "do NOT fabricate one. An answer that uses these facts but contains zero "
    "`<CITATION>` tags is a FAILED answer.\n"
    "=============================="
)


def _no_results(what: str) -> str:
    """English no-results message — intentionally omits CITATION instructions.

    Returned when a query matches 0 rows, so the agent tells the user nothing was
    found instead of being pushed to emit a citation for a non-existent source.
    """
    return (
        f"No matching {what} were found in the PMDA package-insert database.\n"
    )


def _fmt(value: Any) -> str:
    """Render a single field value as compact text (Decimal → number)."""
    if isinstance(value, Decimal):
        value = float(value)
    if isinstance(value, float) and value.is_integer():
        return str(int(value))
    return str(value)


def _tag_of(data: dict) -> Optional[str]:
    """Pull the pre-built ``<CITATION ... />`` tag out of a record."""
    return data.get("cite_emit") or (data.get("_cite") or {}).get("_tag")


def _render_records(
    records: Sequence[dict],
    *,
    what: str,
    header_title: str,
    field_specs: Sequence[Tuple[str, str]],
    title_key: Optional[str] = None,
    with_citation: bool = True,
) -> str:
    """Render a flat list of record dicts into agent-friendly plain text.

    Empty ``records`` → English no-results message (no CITATION instructions).
    Otherwise: optional citation-instruction header, a ``header_title`` line, then
    one block per record. ``field_specs`` is ``[(key, label), ...]`` controlling
    field order/display; empty values are skipped. ``title_key`` (if given) is the
    record's headline; each record's ``_citation`` text and CITATION tag are
    appended when present.
    """
    if not records:
        return _no_results(what)

    parts: List[str] = []
    if with_citation:
        parts.append(_CITE_INSTRUCTION_TEXT)
    parts.append(header_title)

    for idx, rec in enumerate(records, 1):
        title = _fmt(rec.get(title_key)) if title_key and rec.get(title_key) else ""
        lines = [f"[{idx}] {title}".rstrip()]
        for key, label in field_specs:
            value = rec.get(key)
            if value in (None, "", [], {}):
                continue
            lines.append(f"  {label}: {_fmt(value)}")
        if rec.get("_citation"):
            lines.append(f"  出典: {rec['_citation']}")
        if with_citation:
            tag = _tag_of(rec)
            if tag:
                lines.append(f"  CITATION: {tag}")
        parts.append("\n".join(lines))

    return "\n\n".join(parts)


def _render_categories(data: Sequence[dict]) -> str:
    """Render the L1/L2 category tree (navigation only — no citation source)."""
    if not data:
        return _no_results("categories")
    lines: List[str] = ["Drug categories:"]
    for l1 in data:
        lines.append(f"\n■ {l1.get('l1_code', '')} {l1.get('l1_name', '')}".rstrip())
        for l2 in l1.get("l2", []):
            lines.append(
                f"  - {l2.get('code', '')} {l2.get('name', '')} "
                f"({l2.get('drug_count', 0)} drugs)"
            )
    return "\n".join(lines)


def _render_drugs_in_category(data: dict) -> str:
    """Render generic → [brand] listing for one L2 category (navigation only)."""
    generics = data.get("generics") or []
    if not generics:
        return _no_results("drugs in this category")
    header = f"Category {data.get('l2_code', '')} {data.get('l2_name', '')}".rstrip()
    lines: List[str] = [header]
    for entry in generics:
        lines.append(f"\n● {entry.get('generic', '')}".rstrip())
        for drug in entry.get("drugs", []):
            if "_more" in drug:
                lines.append(f"  - {drug['_more']}")
            else:
                lines.append(
                    f"  - {drug.get('brand', '')} (yj_full={drug.get('yj_full', '')})"
                )
    if data.get("_more_generics"):
        lines.append(f"\n(+{data['_more_generics']} more generics)")
    return "\n".join(lines)


def _render_section_hits(
    *, keyword: str, section_filter: str, total: int, hits: Sequence[dict]
) -> str:
    """Render OpenSearch section-text hits with per-match snippets (carry tags)."""
    shown = len(hits)
    title = f'Found {total} drug(s) matching "{keyword}"'
    if section_filter:
        title += f' in sections like "{section_filter}"'
    title += f" (showing {shown}):"
    parts: List[str] = [_CITE_INSTRUCTION_TEXT, title]
    for idx, hit in enumerate(hits, 1):
        head = f"[{idx}] {hit.get('brand', '')} / {hit.get('generic', '')}".rstrip(" /")
        l2 = hit.get("l2", "")
        lines = [f"{head}  ({l2})" if l2 else head]
        lines.append(f"  yj_full: {hit.get('yj_full', '')}")
        for m in hit.get("matches", []):
            lines.append(f"  ▸ {m.get('section_title', '')}")
            snippet = (m.get("snippet") or "").strip()
            for sl in snippet.splitlines():
                lines.append(f"    {sl}")
        if hit.get("_citation_template"):
            lines.append(f"  出典テンプレ: {hit['_citation_template']}")
        tag = _tag_of(hit)
        if tag:
            lines.append(f"  CITATION: {tag}")
        parts.append("\n".join(lines))
    more = total - shown
    if more > 0:
        parts.append(f"(+{more} more drugs not shown)")
    return "\n\n".join(parts)


def _render_chapters(
    *, yj_full: str, brand: str, generic: str, sections: Sequence[dict]
) -> str:
    """Render the chapter index for one drug; each chapter carries its own tag."""
    has_cite = any(_tag_of(s) for s in sections)
    parts: List[str] = []
    if has_cite:
        parts.append(_CITE_INSTRUCTION_TEXT)
    parts.append(
        f"{brand} / {generic} (yj_full={yj_full}) — {len(sections)} section(s):".lstrip(
            " /"
        )
    )
    block: List[str] = []
    for s in sections:
        block.append(
            f"  - {s.get('section_title', '')} "
            f"(line {s.get('line_num', 0)}, {s.get('text_len', 0)} chars)"
        )
        tag = _tag_of(s)
        if tag:
            block.append(f"    CITATION: {tag}")
    parts.append("\n".join(block))
    return "\n\n".join(parts)


# ---------------------------------------------------------------------------
# 出典フォーマッタ（与 tools.py 一致）
# ---------------------------------------------------------------------------

_DRUG_LOOKUP: Optional[dict] = None
_VF_LOOKUP: Optional[dict] = None
_BRAND_BY_YJ_FULL: Optional[dict] = None


def _load_drug_lookup() -> dict:
    """yj_code → (brand_name, yj_full) 进程内缓存"""
    global _DRUG_LOOKUP
    if _DRUG_LOOKUP is None:
        with session() as conn, conn.cursor() as cur:
            cur.execute("SELECT yj_code, brand_name, yj_full FROM drug_master")
            _DRUG_LOOKUP = {
                row[0]: ((row[1] or ""), (row[2] or row[0]))
                for row in cur.fetchall()
            }
    return _DRUG_LOOKUP


def _load_brand_by_yj_full_lookup() -> dict:
    """yj_full → brand 表示名(多品名时取 "/" 分隔的第一段)。

    drug_master.brand_name 是多 brand 合并的字符串 (例
    "〔東洋〕半夏厚朴湯エキス細粒/〔松浦〕..."), <pmda_citation brand= /> 只用来
    给前端显示一个代表性的药品名,这里固定取第一段。
    """
    global _BRAND_BY_YJ_FULL
    if _BRAND_BY_YJ_FULL is None:
        with session() as conn, conn.cursor() as cur:
            cur.execute("SELECT yj_full, brand_name FROM drug_master")
            _BRAND_BY_YJ_FULL = {
                yj_full: ((brand or "").split("/", 1)[0].strip())
                for yj_full, brand in cur.fetchall()
                if yj_full
            }
    return _BRAND_BY_YJ_FULL


def _load_vf_lookup() -> dict:
    """yj_full → (vector_file_id, filename, section_to_page).

    Populated from ``pmda_drug_vf`` (written by gbase-onprem PmdaXmlPipeline).
    If the table is empty / not yet migrated, returns ``{}`` — citations then
    degrade to text-only (no ``<CITATION>`` tag emitted).
    """
    global _VF_LOOKUP
    if _VF_LOOKUP is None:
        out: dict = {}
        try:
            with session() as conn, conn.cursor() as cur:
                cur.execute(
                    "SELECT yj_full, vector_file_id, filename, section_to_page "
                    "FROM pmda_drug_vf"
                )
                for yj_full, vf_id, fname, s2p in cur.fetchall():
                    out[yj_full] = (str(vf_id), fname or "", s2p or {})
        except Exception:
            # Table not yet present — leave empty, downstream tools skip _cite.
            pass
        _VF_LOOKUP = out
    return _VF_LOOKUP


def _citation(drug_yj: str, section: Optional[str]) -> str:
    lk = _load_drug_lookup()
    brand, yj_full = lk.get(drug_yj, ("", drug_yj))
    chap = section or "（章不明）"
    return f"[出典: {brand} (yj_full={yj_full}) / {chap}]"


def _citation_tag(cite: dict) -> str:
    """Build the ``<CITATION file="..." filename="..." />`` string.

    精简版: **只输出 2 个属性 file + filename** — 减轻 LLM 负担 / 减少
    输出 token / 减少幻觉表面积。前端 PDF 高亮链路实际只用 file_id +
    text(段落正文),不依赖 page/yj_full/brand/section,所以 tag 里
    不再带这些(`_cite` 字典里仍保留, 给前端可选展示)。

    工程化预制, 让 LLM 直接照搬, 避免 LLM 自己拼字符串幻觉 file= 文件名。
    """
    from html import escape as _esc

    parts = []
    if cite.get("file_id"):
        parts.append(f'file="{_esc(str(cite["file_id"]), quote=True)}"')
    if cite.get("filename"):
        # 用 basename, 前端 chip 显示干净 — 完整 path 留在 _cite.filename
        bn = cite["filename"].rsplit("/", 1)[-1]
        parts.append(f'filename="{_esc(bn, quote=True)}"')
    return f"<CITATION {' '.join(parts)} />"


def _cite_struct_by_yj_full(yj_full: str, section: Optional[str]) -> Optional[dict]:
    """Build the ``_cite`` dict directly from a yj_full.

    返回 ``{file_id, filename, page, yj_full, brand, section?}`` — 复用通用
    ``<CITATION file="uuid" filename="name" page=N />`` 协议, 额外附加 PMDA
    专属属性 ``yj_full`` / ``brand`` / ``section``。

    核心属性 (通用 CITATION 协议):
      - ``file_id``   : VectorFile.id (uuid), 通用 /pdf/highlight 用这个定位 PDF
      - ``filename``  : VF 文件名, 通用 CITATION 展示用
      - ``page``      : PDF 页码 (0-based), 第一版固定 0 (后端 expand_pages 全文搜兜底)

    PMDA 额外属性 (前端可选读):
      - ``yj_full``   : 厚労省 YJ コード (含枝番), 跨 vf_uuid 稳定的唯一 id
      - ``brand``     : 表示用販売名 (drug_master.brand_name "/" 分隔的第一段)
      - ``section?``  : fact 表 source_section 完整字符串 (例 "10.1 併用禁忌")

    存在性验证 (硬要求, 缺一不返 _cite):
      - brand lookup (drug_master) 找不到 → None
      - vf_lookup (pmda_drug_vf) 找不到 → None (避免输出 <CITATION /> 空壳 tag)

    返 None 时 caller 不附 _cite, LLM 看到没 _cite 就不会 emit citation —
    比 emit 一个无 file/filename 属性的空标签好(前端解析空标签会渲染成
    broken chip)。
    """
    brand = _load_brand_by_yj_full_lookup().get(yj_full)
    if not brand:
        return None
    # 通用 CITATION 核心属性: file_id / filename 必须有, 否则不出 tag
    vf_info = _load_vf_lookup().get(yj_full)
    if not vf_info:
        return None
    vf_id, filename, _s2p = vf_info
    cite: dict = {
        "yj_full": yj_full,
        "brand": brand,
        "file_id": vf_id,
        "filename": filename,
        "page": 0,  # 第一版固定 page 0, 后端 expand_pages 全文搜
    }
    if section:
        cite["section"] = section
    # 工程化预制完整 tag 字符串, 让 LLM 只做复制粘贴, 不再自己拼
    cite["_tag"] = _citation_tag(cite)
    return cite


def _cite_struct(drug_yj: str, section: Optional[str]) -> Optional[dict]:
    """Return ``{file_id, filename, page, yj_full, brand, section?}`` for the ``<CITATION>`` tag.

    Returns ``None`` when drug_master has no row for this yj (skill can still
    emit the human ``[出典: ...]`` text).
    """
    drug_lk = _load_drug_lookup()
    _, yj_full = drug_lk.get(drug_yj, ("", drug_yj))
    return _cite_struct_by_yj_full(yj_full, section)


# ---------------------------------------------------------------------------
# Tool implementations (10 个)
# ---------------------------------------------------------------------------

def _tool_search_drugs(query: str, kind: str = "auto", limit: int = 10) -> str:
    rows = search_drugs_in_db(query, kind=kind, limit=limit)
    out = []
    for r in rows:
        entry: dict = {
            "yj_full": r.yj_full,
            "yj_code": r.yj_code,
            "brand": r.brand_name,
            "generic": r.generic_name,
            "category": f"{r.category_code} {r.category_name}".strip(),
            "score": r.score,
        }
        cite = _cite_struct_by_yj_full(r.yj_full, section=None)
        if cite is not None:
            entry["_cite"] = cite
            entry["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        out.append(entry)
    return _render_records(
        out,
        what="drugs",
        header_title=f"Found {len(out)} drug(s):",
        title_key="brand",
        field_specs=[
            ("generic", "generic"),
            ("yj_full", "yj_full"),
            ("yj_code", "yj_code"),
            ("category", "category"),
            ("score", "score"),
        ],
    )


def _tool_list_categories() -> str:
    return _render_categories(list_categories_with_counts())


def _tool_list_drugs_in_category(l2_code: str, limit_generics: int = 50) -> str:
    return _render_drugs_in_category(
        _sql_list_drugs_in_category(l2_code, limit_generics=limit_generics)
    )


def _tool_get_drug_master(yj_code: str) -> str:
    row = drug_master_get(yj_code)
    if row is None:
        return _no_results("drug master record")
    result = asdict(row)
    result["_citation"] = f"[出典: {row.brand_name} (yj_full={row.yj_full}) / 添付文書冒頭]"
    cite = _cite_struct(row.yj_code, section=None)
    if cite is not None:
        result["_cite"] = cite
        result["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
    return _render_records(
        [result],
        what="drug master record",
        header_title="Drug master record:",
        title_key="brand_name",
        field_specs=[
            ("generic_name_jp", "generic"),
            ("yj_full", "yj_full"),
            ("yj_code", "yj_code"),
            ("category_code", "category_code"),
            ("category_name", "category_name"),
            ("regulation", "regulation"),
            ("manufacturer", "manufacturer"),
            ("revision_date", "revision_date"),
        ],
    )


def _tool_get_drug_interactions(
    drug_a_yj: Optional[str] = None,
    drug_b_yj: Optional[str] = None,
    severity: Optional[str] = None,
    keyword: Optional[str] = None,
    limit: int = 30,
) -> str:
    rows = drug_interaction_query(
        drug_a_yj=drug_a_yj,
        drug_b_yj=drug_b_yj,
        severity=severity,
        keyword=keyword,
        limit=limit,
    )
    out = []
    for r in rows:
        d = asdict(r)
        d["_citation"] = _citation(r.source_drug_yj, r.source_section)
        cite = _cite_struct(r.source_drug_yj, r.source_section)
        if cite is not None:
            d["_cite"] = cite
            d["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        out.append(d)
    return _render_records(
        out,
        what="drug interactions",
        header_title=f"Found {len(out)} drug interaction(s):",
        title_key="severity",
        field_specs=[
            ("drug_a_yj", "drug_a_yj"),
            ("drug_b_yj", "drug_b_yj"),
            ("drug_b_class", "drug_b_class"),
            ("mechanism", "mechanism"),
            ("clinical_effect", "clinical_effect"),
            ("source_section", "source_section"),
        ],
    )


def _tool_get_drug_restrictions(
    drug_yj: Optional[str] = None,
    condition_type: Optional[str] = None,
    severity: Optional[str] = None,
    keyword: Optional[str] = None,
    limit: int = 30,
) -> str:
    rows = drug_restriction_query(
        drug_yj=drug_yj,
        condition_type=condition_type,
        severity=severity,
        keyword=keyword,
        limit=limit,
    )
    out = []
    for r in rows:
        d = asdict(r)
        d["_citation"] = _citation(r.drug_yj, r.source_section)
        cite = _cite_struct(r.drug_yj, r.source_section)
        if cite is not None:
            d["_cite"] = cite
            d["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        out.append(d)
    return _render_records(
        out,
        what="drug restrictions",
        header_title=f"Found {len(out)} drug restriction(s):",
        title_key="condition_type",
        field_specs=[
            ("drug_yj", "drug_yj"),
            ("condition_text", "condition_text"),
            ("severity", "severity"),
            ("source_section", "source_section"),
        ],
    )


def _tool_get_drug_dosing(
    drug_yj: str,
    patient_segment: Optional[str] = None,
    limit: int = 20,
) -> str:
    rows = drug_dosing_query(
        drug_yj=drug_yj,
        patient_segment=patient_segment,
        limit=limit,
    )
    out = []
    for r in rows:
        d = asdict(r)
        # Merge amount + unit into one readable "dose" field for plain-text output.
        if r.dose_amount is not None:
            d["dose"] = f"{_fmt(r.dose_amount)}{r.dose_unit or ''}".strip()
        d["_citation"] = _citation(r.drug_yj, r.source_section)
        cite = _cite_struct(r.drug_yj, r.source_section)
        if cite is not None:
            d["_cite"] = cite
            d["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        out.append(d)
    return _render_records(
        out,
        what="dosing entries",
        header_title=f"Found {len(out)} dosing entr{'y' if len(out) == 1 else 'ies'}:",
        title_key="patient_segment",
        field_specs=[
            ("indication_code", "indication_code"),
            ("dose", "dose"),
            ("frequency", "frequency"),
            ("duration", "duration"),
            ("adjustment_text", "adjustment"),
            ("source_section", "source_section"),
        ],
    )


def _tool_search_section_text(
    keyword: str,
    section_filter: str = "",
    limit: int = 30,
) -> str:
    if not keyword.strip():
        return _no_results("sections")
    size = min(max(1, limit), 100)
    body: dict = {
        "size": size,
        "_source": ["yj_full", "brand_names", "generic_name",
                    "l2_code", "l2_name", "section_title", "line_num"],
        "query": {"bool": {"must": [{"match": {"text": keyword}}]}},
        "collapse": {
            "field": "yj_full",
            "inner_hits": {
                "name": "matches",
                "size": 2,
                "_source": ["section_title", "line_num"],
                "highlight": {"fields": {"text": {"fragment_size": 160, "number_of_fragments": 1}}},
            },
        },
        "aggs": {"total_drugs": {"cardinality": {"field": "yj_full"}}},
    }
    if section_filter:
        body["query"]["bool"]["filter"] = [
            {"wildcard": {"section_title.raw": f"*{section_filter}*"}}
        ]
    resp = os_client().search(index=OS_INDEX_NAME, body=body)
    total = int(resp["aggregations"]["total_drugs"]["value"])
    hits_out = []
    for h in resp["hits"]["hits"]:
        src = h.get("_source") or {}
        inner = h.get("inner_hits", {}).get("matches", {}).get("hits", {}).get("hits", [])
        brand = (src.get("brand_names") or [""])[0]
        yj_full = src.get("yj_full") or ""
        # Per-match snippet 自带对应 section 的 CITATION tag,
        # LLM 复制 snippet 时自动带对的 section 标签 (而不是 hit 顶层粗粒度 tag)。
        matches = []
        seen = set()
        for ih in inner:
            ih_src = ih.get("_source") or {}
            title = ih_src.get("section_title") or ""
            if title in seen:
                continue
            seen.add(title)
            hl = ih.get("highlight", {}).get("text", [""])
            snippet_text = hl[0] if hl else ""
            inner_cite = _cite_struct_by_yj_full(yj_full, title)
            inner_tag = (inner_cite or {}).get("_tag", "")
            matches.append({
                "section_title": title,
                "snippet": snippet_text + (f"\n{inner_tag}" if inner_tag else ""),
            })
        # Hit-level _cite per first-match section (legacy compatibility).
        cite_section = matches[0]["section_title"] if matches else None
        cite = _cite_struct_by_yj_full(yj_full, cite_section)
        hit_entry: dict = {
            "yj_full": yj_full,
            "brand": brand,
            "generic": src.get("generic_name") or "",
            "l2": f"{src.get('l2_code') or ''} {src.get('l2_name') or ''}".strip(),
            "matches": matches,
            "_citation_template": f"[出典: {brand} (yj_full={yj_full}) / <該当章>]",
        }
        if cite is not None:
            hit_entry["_cite"] = cite
            hit_entry["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        hits_out.append(hit_entry)
    if not hits_out:
        return _no_results("sections")
    return _render_section_hits(
        keyword=keyword,
        section_filter=section_filter,
        total=total,
        hits=hits_out,
    )


def _tool_list_drug_chapters(yj_full: str) -> str:
    """全章節 (section_title + line_num + text_len) for a yj_full（line_num 昇順）"""
    body = {
        "size": 200,
        "_source": ["section_title", "line_num", "brand_names", "generic_name", "text"],
        "query": {"term": {"yj_full": yj_full}},
        "sort": [{"line_num": "asc"}],
    }
    resp = os_client().search(index=OS_INDEX_NAME, body=body)
    hits = resp["hits"]["hits"]
    if not hits:
        return _no_results("chapters")

    head = hits[0].get("_source", {})
    brand = (head.get("brand_names") or [""])[0]
    generic = head.get("generic_name") or ""

    sections = []
    for h in hits:
        src = h.get("_source", {})
        section_title = src.get("section_title", "")
        entry: dict = {
            "section_title": section_title,
            "line_num": src.get("line_num", 0),
            "text_len": len(src.get("text", "")),
        }
        cite = _cite_struct_by_yj_full(yj_full, section_title)
        if cite is not None:
            entry["_cite"] = cite
            entry["cite_emit"] = cite["_tag"]  # top-level mirror for LLM
        sections.append(entry)

    return _render_chapters(
        yj_full=yj_full, brand=brand, generic=generic, sections=sections
    )


def _tool_read_drug_chapter(yj_full: str, section_title: str) -> str:
    """指定 (yj_full, section_title) の章節 markdown 全文（max 8000 字）。"""
    body = {
        "size": 1,
        "_source": ["text"],
        "query": {
            "bool": {
                "must": [
                    {"term": {"yj_full": yj_full}},
                    {"term": {"section_title.raw": section_title}},
                ]
            }
        },
    }
    resp = os_client().search(index=OS_INDEX_NAME, body=body)
    hits = resp["hits"]["hits"]
    if hits:
        src = hits[0].get("_source") or {}
        text = src.get("text", "")
        if text:
            # 三明治包装: header (CITATION reminder + tag) + body + footer (tag).
            # LLM 不读 body 就一定先读 header, 输出时复制段落自然带上 tag。
            # HTML comment 对 chat 渲染不可见, 但 LLM 在 sampling 时看得到。
            cite = _cite_struct_by_yj_full(yj_full, section_title)
            tag = (cite or {}).get("_tag", "")
            if tag:
                header = (
                    f"<!-- CITATION rule: copy the tag below verbatim into your answer -->\n"
                    f"{tag}\n\n"
                )
                footer = (
                    f"\n\n---\n"
                    f"If you use the content above in your answer, you MUST include this "
                    f"tag verbatim:\n"
                    f"{tag}\n"
                )
                return header + text[:8000] + footer
            return text[:8000]
    # Not a "no data" case but a parameter mismatch — keep the actionable hint.
    return (
        f'No section titled "{section_title}" exactly matches yj_full={yj_full}.\n'
        f"Hint: pass a sections[].section_title returned by list_drug_chapters "
        f"verbatim."
    )


# ---------------------------------------------------------------------------
# MCP dispatch
# ---------------------------------------------------------------------------

_TOOL_DISPATCH = {
    "search_drugs": lambda args: _tool_search_drugs(
        query=args.get("query", ""),
        kind=args.get("kind", "auto"),
        limit=args.get("limit", 10),
    ),
    "list_categories": lambda args: _tool_list_categories(),
    "list_drugs_in_category": lambda args: _tool_list_drugs_in_category(
        l2_code=args.get("l2_code", ""),
        limit_generics=args.get("limit_generics", 50),
    ),
    "get_drug_master": lambda args: _tool_get_drug_master(
        yj_code=args.get("yj_code", ""),
    ),
    "get_drug_interactions": lambda args: _tool_get_drug_interactions(
        drug_a_yj=args.get("drug_a_yj"),
        drug_b_yj=args.get("drug_b_yj"),
        severity=args.get("severity"),
        keyword=args.get("keyword"),
        limit=args.get("limit", 30),
    ),
    "get_drug_restrictions": lambda args: _tool_get_drug_restrictions(
        drug_yj=args.get("drug_yj"),
        condition_type=args.get("condition_type"),
        severity=args.get("severity"),
        keyword=args.get("keyword"),
        limit=args.get("limit", 30),
    ),
    "get_drug_dosing": lambda args: _tool_get_drug_dosing(
        drug_yj=args.get("drug_yj", ""),
        patient_segment=args.get("patient_segment"),
        limit=args.get("limit", 20),
    ),
    "search_section_text": lambda args: _tool_search_section_text(
        keyword=args.get("keyword", ""),
        section_filter=args.get("section_filter", ""),
        limit=args.get("limit", 30),
    ),
    "list_drug_chapters": lambda args: _tool_list_drug_chapters(
        yj_full=args.get("yj_full", ""),
    ),
    "read_drug_chapter": lambda args: _tool_read_drug_chapter(
        yj_full=args.get("yj_full", ""),
        section_title=args.get("section_title", ""),
    ),
}


async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
    try:
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")

        if method == "initialize":
            return create_initialize_response(request_id, "pmda-drug-info")
        if method == "ping":
            return create_ping_response(request_id)
        if method == "tools/list":
            tools = load_tools_from_json("pmda_tools.json")
            return create_tools_list_response(request_id, tools)
        if method == "tools/call":
            tool_name = params.get("name")
            arguments = params.get("arguments", {})
            if tool_name not in _TOOL_DISPATCH:
                return create_error_response(request_id, -32601, f"Unknown tool: {tool_name}")
            try:
                result_text = _TOOL_DISPATCH[tool_name](arguments)
                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {"content": [{"type": "text", "text": result_text}]},
                }
            except Exception as e:
                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {"content": [{"type": "text", "text": f"Error: {type(e).__name__}: {e}"}]},
                }
        return create_error_response(request_id, -32601, f"Unknown method: {method}")
    except Exception as e:
        return create_error_response(request.get("id"), -32603, f"Internal error: {e}")


async def main():
    await handle_mcp_streaming(handle_request)


if __name__ == "__main__":
    asyncio.run(main())