#!/usr/bin/env python3
"""
PMDA drug information MCP server (mock data version).

Provides drug search, master info, interactions, restrictions, dosing,
and full-text chapter retrieval with mock data for testing.
"""

import asyncio
import json
import sys
from typing import Any, Dict, Optional

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


def _dump(obj) -> str:
    return json.dumps(obj, ensure_ascii=False)


# ---------------------------------------------------------------------------
# Mock data
# ---------------------------------------------------------------------------

MOCK_DRUG_MASTER = {
    "2149039F1082": {
        "yj_code": "2149039F1082",
        "yj_full": "2149039F1082_1_17",
        "brand_name": "ロサルタンK錠50mg「科研」",
        "generic_name": "ロサルタンカリウム",
        "category_code": "214",
        "category_name": "アンジオテンシンII受容体拮抗薬",
        "regulation": "劇薬, 処方箋医薬品",
        "manufacturer": "科研製薬株式会社",
        "revision_date": "2024-06",
    },
    "3399007H1021": {
        "yj_code": "3399007H1021",
        "yj_full": "3399007H1021_1_21",
        "brand_name": "バイアスピリン錠100mg",
        "generic_name": "アスピリン",
        "category_code": "339",
        "category_name": "血液・体液用薬",
        "regulation": "処方箋医薬品",
        "manufacturer": "バイエル薬品株式会社",
        "revision_date": "2024-03",
    },
    "2179004F1026": {
        "yj_code": "2179004F1026",
        "yj_full": "2179004F1026_1_14",
        "brand_name": "ノルバスク錠5mg",
        "generic_name": "アムロジピンベシル酸塩",
        "category_code": "217",
        "category_name": "カルシウム拮抗薬",
        "regulation": "処方箋医薬品",
        "manufacturer": "ファイザー株式会社",
        "revision_date": "2024-01",
    },
}

MOCK_CATEGORIES = [
    {"category_code": "214", "category_name": "アンジオテンシンII受容体拮抗薬", "level": "L2", "drug_count": 35},
    {"category_code": "217", "category_name": "カルシウム拮抗薬", "level": "L2", "drug_count": 48},
    {"category_code": "339", "category_name": "血液・体液用薬", "level": "L2", "drug_count": 22},
    {"category_code": "612", "category_name": "消化性潰瘍用剤", "level": "L2", "drug_count": 40},
]

MOCK_INTERACTIONS = [
    {
        "drug_a_yj": "2149039F1082",
        "drug_b_yj": "3399007H1021",
        "drug_b_class": "アスピリン（抗血小板剤）",
        "severity": "併用注意",
        "mechanism": "ARBの降圧作用を減弱するおそれがある。また、腎機能低下・高カリウム血症のリスクを増大。",
        "clinical_effect": "降圧効果の減弱、腎機能悪化、高カリウム血症に注意。",
        "source_drug_yj": "2149039F1082",
        "source_section": "10.2 併用注意",
    },
    {
        "drug_a_yj": "3399007H1021",
        "drug_b_yj": "2149039F1082",
        "drug_b_class": "ロサルタンカリウム（ARB）",
        "severity": "併用注意",
        "mechanism": "アスピリンの副作用（消化性潰瘍、腎機能低下）を増強するおそれ。",
        "clinical_effect": "消化性潰瘍、腎機能低下に注意。血清カリウム値の上昇に注意。",
        "source_drug_yj": "3399007H1021",
        "source_section": "10.2 併用注意",
    },
]

MOCK_RESTRICTIONS = [
    {
        "drug_yj": "2149039F1082",
        "condition_type": "腎機能障害",
        "condition_text": "腎機能障害患者",
        "condition_params": {"eGFR_max": 30},
        "severity": "慎重投与",
        "source_section": "9.2 腎機能障害患者",
    },
    {
        "drug_yj": "2149039F1082",
        "condition_type": "妊婦",
        "condition_text": "妊娠中の女性",
        "condition_params": {},
        "severity": "禁忌",
        "source_section": "9.5 妊婦",
    },
    {
        "drug_yj": "2149039F1082",
        "condition_type": "高齢者",
        "condition_text": "高齢者（65歳以上）",
        "condition_params": {},
        "severity": "慎重投与",
        "source_section": "9.8 高齢者",
    },
    {
        "drug_yj": "3399007H1021",
        "condition_type": "過敏症",
        "condition_text": "本剤の成分に対し過敏症の既往歴のある患者",
        "condition_params": {},
        "severity": "禁忌",
        "source_section": "2. 禁忌",
    },
]

MOCK_DOSING = [
    {
        "drug_yj": "2149039F1082",
        "patient_segment": "成人",
        "segment_params": {},
        "indication_code": "高血圧症",
        "dose_amount": "50",
        "dose_unit": "mg",
        "frequency": "1日1回",
        "duration": "",
        "adjustment_text": "効果不十分な場合は100mgまで増量可",
        "source_section": "6. 用法及び用量",
    },
    {
        "drug_yj": "2149039F1082",
        "patient_segment": "腎機能障害患者",
        "segment_params": {"eGFR_max": 30},
        "indication_code": "高血圧症",
        "dose_amount": "25",
        "dose_unit": "mg",
        "frequency": "1日1回",
        "duration": "",
        "adjustment_text": "eGFR 30以下では用量を減ずること。血清カリウム・クレアチニンの推移に注意。",
        "source_section": "9.2 腎機能障害患者",
    },
]

MOCK_CHAPTERS = {
    "2149039F1082_1_17": [
        {"section_title": "1. 警告", "line_num": 1, "text_len": 120},
        {"section_title": "2. 禁忌", "line_num": 5, "text_len": 80},
        {"section_title": "4. 効能・効果", "line_num": 12, "text_len": 60},
        {"section_title": "6. 用法及び用量", "line_num": 20, "text_len": 150},
        {"section_title": "9.2 腎機能障害患者", "line_num": 45, "text_len": 200},
        {"section_title": "9.5 妊婦", "line_num": 52, "text_len": 180},
        {"section_title": "9.8 高齢者", "line_num": 60, "text_len": 100},
        {"section_title": "10.2 併用注意", "line_num": 75, "text_len": 350},
        {"section_title": "11.1 重大な副作用", "line_num": 90, "text_len": 400},
        {"section_title": "11.2 その他の副作用", "line_num": 110, "text_len": 300},
    ],
    "3399007H1021_1_21": [
        {"section_title": "1. 警告", "line_num": 1, "text_len": 100},
        {"section_title": "2. 禁忌", "line_num": 4, "text_len": 90},
        {"section_title": "4. 効能・効果", "line_num": 10, "text_len": 55},
        {"section_title": "6. 用法及び用量", "line_num": 18, "text_len": 130},
        {"section_title": "10.2 併用注意", "line_num": 70, "text_len": 300},
        {"section_title": "11.1 重大な副作用", "line_num": 85, "text_len": 450},
        {"section_title": "11.2 その他の副作用", "line_num": 105, "text_len": 280},
    ],
}

MOCK_SECTION_TEXT = {
    ("2149039F1082_1_17", "9.2 腎機能障害患者"): (
        "9.2 腎機能障害患者\n"
        "腎機能障害患者（eGFR 30 mL/min/1.73m²以下）には、ロサルタンカリウムの"
        "投与開始用量を25mg/日とし、血清カリウム及び血清クレアチニンの推移に"
        "十分注意すること。\n"
        "【理由】腎機能障害患者では、本剤の投与により急速に腎機能が悪化する"
        "おそれがある。また、高カリウム血症があらわれやすい。"
    ),
    ("2149039F1082_1_17", "9.5 妊婦"): (
        "9.5 妊婦\n"
        "妊婦又は妊娠している可能性のある女性には投与しないこと。\n"
        "【理由】妊娠中期・末期にレニン-アンジオテンシン系に作用する薬剤を"
        "投与された患者では、胎児の腎機能低下、羊水過少症、頭蓋の発育不全、"
        "肺低形成等があらわれるおそれがある。"
    ),
    ("2149039F1082_1_17", "10.2 併用注意"): (
        "10.2 併用注意\n"
        "・アスピリン（抗血小板剤）\n"
        "  【リスク】ARBの降圧作用を減弱するおそれがある。\n"
        "  腎機能低下・高カリウム血症のリスクを増大。\n"
        "  【措置】降圧効果の減弱、腎機能悪化、高カリウム血症に注意すること。"
    ),
    ("2149039F1082_1_17", "11.1 重大な副作用"): (
        "11.1 重大な副作用\n"
        "・血管浮腫（頻度不明）：顔面、口唇、咽頭、舌等の腫脹があらわれた場合には"
        "直ちに投与を中止し、適切な処置を行うこと。\n"
        "・高カリウム血症（0.1%未満）：血清カリウム値の上昇があらわれることがある。\n"
        "・腎機能悪化（0.1%未満）：BUN、クレアチニンの上昇があらわれることがある。"
    ),
    ("3399007H1021_1_21", "10.2 併用注意"): (
        "10.2 併用注意\n"
        "・ロサルタンカリウム（ARB）\n"
        "  【リスク】アスピリンの副作用（消化性潰瘍、腎機能低下）を増強するおそれ。\n"
        "  【措置】消化性潰瘍、腎機能低下に注意。血清カリウム値の上昇に注意すること。"
    ),
    ("3399007H1021_1_21", "11.1 重大な副作用"): (
        "11.1 重大な副作用\n"
        "・ショック、アナフィラキシー（頻度不明）：呼吸困難、血圧低下等があらわれた\n"
        "  場合には直ちに投与を中止し、適切な処置を行うこと。\n"
        "・消化性潰瘍（0.1%未満）：出血、穿孔があらわれることがある。\n"
        "・腎機能障害（0.1%未満）：急性腎不全があらわれることがある。"
    ),
}


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


# ---------------------------------------------------------------------------
# Tool implementations (mock)
# ---------------------------------------------------------------------------

def _tool_search_drugs(query: str, kind: str = "auto", limit: int = 10) -> str:
    results = []
    for code, d in MOCK_DRUG_MASTER.items():
        q = query.lower()
        if (kind == "brand" and q in d["brand_name"].lower()) or \
           (kind == "generic" and q in d["generic_name"].lower()) or \
           (kind == "yj" and (q in d["yj_code"].lower() or q in d["yj_full"].lower())) or \
           (kind == "auto" and (q in d["brand_name"].lower() or q in d["generic_name"].lower()
                                or q in d["yj_code"].lower() or q in d["yj_full"].lower())):
            results.append({
                "yj_full": d["yj_full"],
                "yj_code": d["yj_code"],
                "brand": d["brand_name"],
                "generic": d["generic_name"],
                "category": f"{d['category_code']} {d['category_name']}",
                "score": 1.0,
            })
    return _dump(results[:limit])


def _tool_list_categories() -> str:
    return _dump(MOCK_CATEGORIES)


def _tool_list_drugs_in_category(l2_code: str, limit_generics: int = 50) -> str:
    results = []
    seen_generics = set()
    for code, d in MOCK_DRUG_MASTER.items():
        if d["category_code"].startswith(l2_code) and d["generic_name"] not in seen_generics:
            seen_generics.add(d["generic_name"])
            results.append({
                "generic_name": d["generic_name"],
                "brands": [{"yj_code": d["yj_code"], "brand_name": d["brand_name"], "yj_full": d["yj_full"]}],
            })
    return _dump(results[:limit_generics])


def _tool_get_drug_master(yj_code: str) -> str:
    d = MOCK_DRUG_MASTER.get(yj_code)
    if not d:
        return _dump({"error": f"yj_code {yj_code} not found"})
    result = dict(d)
    result["_citation"] = f"[出典: {d['brand_name']} (yj_full={d['yj_full']}) / 添付文書冒頭]"
    return _dump(result)


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:
    results = []
    for r in MOCK_INTERACTIONS:
        if drug_a_yj and r["drug_a_yj"] != drug_a_yj:
            continue
        if drug_b_yj and r["drug_b_yj"] != drug_b_yj:
            continue
        if severity and r["severity"] != severity:
            continue
        if keyword and keyword.lower() not in (
            (r.get("drug_b_class") or "").lower()
            + (r.get("mechanism") or "").lower()
            + (r.get("clinical_effect") or "").lower()
        ):
            continue
        results.append({**r, "_citation": _citation(r["source_drug_yj"], r["source_section"])})
    return _dump(results[:limit])


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:
    results = []
    for r in MOCK_RESTRICTIONS:
        if drug_yj and r["drug_yj"] != drug_yj:
            continue
        if condition_type and r["condition_type"] != condition_type:
            continue
        if severity and r["severity"] != severity:
            continue
        if keyword and keyword.lower() not in (r.get("condition_text") or "").lower():
            continue
        results.append({**r, "_citation": _citation(r["drug_yj"], r["source_section"])})
    return _dump(results[:limit])


def _tool_get_drug_dosing(
    drug_yj: str,
    patient_segment: Optional[str] = None,
    limit: int = 20,
) -> str:
    results = []
    for r in MOCK_DOSING:
        if r["drug_yj"] != drug_yj:
            continue
        if patient_segment and r["patient_segment"] != patient_segment:
            continue
        results.append({**r, "_citation": _citation(drug_yj, r["source_section"])})
    return _dump(results[:limit])


def _tool_search_section_text(
    keyword: str,
    section_filter: str = "",
    limit: int = 30,
) -> str:
    if not keyword.strip():
        return _dump({"keyword": keyword, "total_drugs": 0, "shown": 0, "hits": []})

    # Simple mock: search through section text
    hits_out = []
    for (yj_full, section_title), text in MOCK_SECTION_TEXT.items():
        if section_filter and section_filter not in section_title:
            continue
        if keyword.lower() in text.lower():
            drug = None
            for d in MOCK_DRUG_MASTER.values():
                if d["yj_full"] == yj_full:
                    drug = d
                    break
            if not drug:
                continue
            brand = drug["brand_name"]
            # Deduplicate by yj_full
            existing = [h for h in hits_out if h["yj_full"] == yj_full]
            if existing:
                existing[0]["matches"].append({
                    "section_title": section_title,
                    "snippet": text[:160],
                })
                continue
            hits_out.append({
                "yj_full": yj_full,
                "brand": brand,
                "generic": drug["generic_name"],
                "l2": f"{drug['category_code']} {drug['category_name']}",
                "matches": [{"section_title": section_title, "snippet": text[:160]}],
                "_citation_template": f"[出典: {brand} (yj_full={yj_full}) / <該当章>]",
            })

    return _dump({
        "keyword": keyword,
        "section_filter": section_filter or None,
        "total_drugs": len({h["yj_full"] for h in hits_out}),
        "shown": len(hits_out),
        "hits": hits_out[:limit],
    })


def _tool_list_drug_chapters(yj_full: str) -> str:
    sections = MOCK_CHAPTERS.get(yj_full)
    if not sections:
        return _dump({"error": f"yj_full {yj_full} の章節が見つかりません。"})

    drug = None
    for d in MOCK_DRUG_MASTER.values():
        if d["yj_full"] == yj_full:
            drug = d
            break

    return _dump({
        "yj_full": yj_full,
        "brand": drug["brand_name"] if drug else "",
        "generic": drug["generic_name"] if drug else "",
        "n_sections": len(sections),
        "sections": sections,
    })


def _tool_read_drug_chapter(yj_full: str, section_title: str) -> str:
    text = MOCK_SECTION_TEXT.get((yj_full, section_title))
    if text:
        return text[:8000]
    return _dump({
        "error": f"section_title {section_title!r} は {yj_full} に存在しません。",
        "hint": "list_drug_chapters で取得した sections[].section_title をそのまま渡してください。",
    })


# ---------------------------------------------------------------------------
# MCP request handler
# ---------------------------------------------------------------------------

_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]:
    """Handle an MCP request."""
    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")

        elif method == "ping":
            return create_ping_response(request_id)

        elif method == "tools/list":
            tools = load_tools_from_json("pmda_tools.json")
            return create_tools_list_response(request_id, tools)

        elif 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: {str(e)}"}]
                    },
                }

        else:
            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: {str(e)}")


async def main():
    await handle_mcp_streaming(handle_request)


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