"""Demo 7 — Medical Triage Assistant with Self-Organizing Graph.

A graph-augmented medical triage assistant that helps with patient lookup,
vitals review, lab ordering, and medication checks.  Policies emerge from
repeated usage patterns — no hardcoded rules.

Requirements:
  - FalkorDB running:  docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d
  - LLM API key in .env (OPENAI_API_KEY or ANTHROPIC_API_KEY)

Usage:
  uv run python demos/demo7_medical.py
  Open http://localhost:8007
"""

from __future__ import annotations

import json
import re
import sys
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any

sys.path.insert(0, str(Path(__file__).resolve().parent))

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from ui import chat_page, info_row

from koan.agents.events import AgentComplete, AgentError, TextDelta, ToolCallEnd, ToolCallStart
from koan.agents.types import AgentConfig
from koan.graph import (
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    create_graph_client,
    record_feedback,
)
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.tools import ToolRegistry, tool

if TYPE_CHECKING:
    from koan.graph.client import GraphClient

load_dotenv()

# ---------------------------------------------------------------------------
# Stub data — patients, vitals, labs, medications
# ---------------------------------------------------------------------------

PATIENTS: dict[str, dict[str, Any]] = {
    "pt-4001": {
        "name": "Maria Garcia",
        "age": 45,
        "sex": "F",
        "blood_type": "O+",
        "allergies": ["penicillin"],
        "conditions": ["hypertension", "type 2 diabetes"],
        "primary_care": "Dr. Nguyen",
    },
    "pt-4002": {
        "name": "James Wilson",
        "age": 67,
        "sex": "M",
        "blood_type": "A-",
        "allergies": [],
        "conditions": ["atrial fibrillation", "COPD"],
        "primary_care": "Dr. Patel",
    },
    "pt-4003": {
        "name": "Sarah Chen",
        "age": 32,
        "sex": "F",
        "blood_type": "B+",
        "allergies": ["sulfa drugs", "latex"],
        "conditions": [],
        "primary_care": "Dr. Thompson",
    },
}

VITALS: dict[str, dict[str, Any]] = {
    "pt-4001": {
        "bp": "148/92",
        "hr": 82,
        "temp": 98.6,
        "spo2": 97,
        "weight_kg": 78.5,
        "recorded": "2026-03-10 08:30",
        "flags": ["BP elevated — above target for diabetic patient"],
    },
    "pt-4002": {
        "bp": "130/80",
        "hr": 110,
        "temp": 99.1,
        "spo2": 91,
        "weight_kg": 85.0,
        "recorded": "2026-03-10 07:15",
        "flags": ["HR elevated — check rhythm", "SpO2 low — consider supplemental O2"],
    },
    "pt-4003": {
        "bp": "118/72",
        "hr": 74,
        "temp": 101.3,
        "spo2": 98,
        "weight_kg": 62.0,
        "recorded": "2026-03-10 09:00",
        "flags": ["Temp elevated — possible infection"],
    },
}

MEDICATIONS: dict[str, list[dict[str, str]]] = {
    "pt-4001": [
        {"name": "Metformin", "dose": "500mg", "frequency": "twice daily", "for": "diabetes"},
        {"name": "Lisinopril", "dose": "10mg", "frequency": "once daily", "for": "hypertension"},
    ],
    "pt-4002": [
        {"name": "Warfarin", "dose": "5mg", "frequency": "once daily", "for": "atrial fibrillation"},
        {"name": "Albuterol", "dose": "90mcg", "frequency": "as needed", "for": "COPD"},
        {"name": "Tiotropium", "dose": "18mcg", "frequency": "once daily", "for": "COPD"},
    ],
    "pt-4003": [],
}


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@tool
def lookup_patient(patient_id: str) -> str:
    """Look up patient demographics and medical history.

    Args:
        patient_id: The patient identifier, e.g. 'pt-4001'.
    """
    pt = PATIENTS.get(patient_id)
    if pt is None:
        return f"Patient {patient_id} not found."
    return json.dumps({"patient_id": patient_id, **pt})


@tool
def check_vitals(patient_id: str) -> str:
    """Check the latest vitals for a patient including any flags.

    Args:
        patient_id: The patient identifier.
    """
    vitals = VITALS.get(patient_id)
    if vitals is None:
        return f"No vitals recorded for {patient_id}."
    return json.dumps({"patient_id": patient_id, **vitals})


@tool
def order_labs(patient_id: str, tests: str) -> str:
    """Order laboratory tests for a patient.

    Args:
        patient_id: The patient identifier.
        tests: Comma-separated list of tests to order (e.g. 'CBC, BMP, HbA1c').
    """
    pt = PATIENTS.get(patient_id)
    if pt is None:
        return f"Patient {patient_id} not found."
    test_list = [t.strip() for t in tests.split(",")]
    return (
        f"LAB ORDER — {pt['name']} ({patient_id})\n"
        f"Tests ordered: {', '.join(test_list)}\n"
        f"Priority: Routine\n"
        f"Status: Order submitted. Results expected in 2-4 hours."
    )


@tool
def check_medications(patient_id: str) -> str:
    """Review current medications for a patient.

    Args:
        patient_id: The patient identifier.
    """
    meds = MEDICATIONS.get(patient_id)
    if meds is None:
        return f"Patient {patient_id} not found."
    if not meds:
        return f"No current medications for {patient_id}."
    return json.dumps({"patient_id": patient_id, "medications": meds, "count": len(meds)})


@tool
def flag_allergy_risk(patient_id: str, proposed_medication: str) -> str:
    """Check if a proposed medication conflicts with patient allergies.

    Args:
        patient_id: The patient identifier.
        proposed_medication: The medication being considered.
    """
    pt = PATIENTS.get(patient_id)
    if pt is None:
        return f"Patient {patient_id} not found."

    allergies = pt.get("allergies", [])
    med_lower = proposed_medication.lower()

    # Simple allergy cross-reference
    penicillin_family = ["penicillin", "amoxicillin", "ampicillin", "augmentin"]
    sulfa_family = ["sulfamethoxazole", "bactrim", "septra", "sulfa"]

    for allergy in allergies:
        if allergy.lower() == "penicillin" and med_lower in penicillin_family:
            return f"⚠️ ALLERGY ALERT: {pt['name']} is allergic to {allergy}. {proposed_medication} is contraindicated."
        if allergy.lower() == "sulfa drugs" and med_lower in sulfa_family:
            return f"⚠️ ALLERGY ALERT: {pt['name']} is allergic to {allergy}. {proposed_medication} is contraindicated."

    return f"No allergy conflict found for {proposed_medication} in patient {pt['name']}."


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(title="Demo 7: Medical Triage")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

registry = ToolRegistry()
registry.register(lookup_patient)
registry.register(check_vitals)
registry.register(order_labs)
registry.register(check_medications)
registry.register(flag_allergy_risk)

graph_config = GraphConfig(
    provider="falkordb",
    uri="redis://localhost:6379",
    database="koan_demo7",
)

graph_client: GraphClient | None = None

agent_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.2,
    system_prompt=(
        "You are a medical triage assistant. Help clinicians with patient lookups, "
        "vitals review, lab ordering, medication checks, and allergy screening. "
        "Always flag critical vitals and allergy risks. Be precise with medical data. "
        "Use the provided patient ID directly — do not ask for confirmation."
    ),
)

threads: dict[str, list[dict[str, str]]] = {}


@app.on_event("startup")
async def startup() -> None:
    global graph_client  # noqa: PLW0603
    graph_client = create_graph_client(graph_config)
    await graph_client.connect()


@app.on_event("shutdown")
async def shutdown() -> None:
    if graph_client is not None:
        await graph_client.close()


# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------


class ChatRequest(BaseModel):
    thread_id: str
    message: str


class FeedbackRequest(BaseModel):
    run_id: str
    satisfied: bool


class CreateThreadResponse(BaseModel):
    thread_id: str


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(content=chat_page(
        title="Demo 7: Medical Triage",
        badge_text="Medical",
        badge_color="red",
        port=8007,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Tools", "lookup_patient, check_vitals, order_labs, check_medications, flag_allergy_risk")
            + info_row("Domain", "Medical / Triage")
            + info_row("Graph", "FalkorDB (koan_demo7)")
        ),
    ))


@app.get("/api/threads")
async def list_threads() -> list[dict[str, str]]:
    return [{"thread_id": tid} for tid in threads]


@app.post("/api/threads")
async def create_thread() -> CreateThreadResponse:
    thread_id = str(uuid.uuid4())
    threads[thread_id] = []
    return CreateThreadResponse(thread_id=thread_id)


@app.get("/api/threads/{thread_id}/messages")
async def get_messages(thread_id: str) -> list[dict]:
    return [
        {"role": m["role"], "content": m["content"], "metadata": m.get("metadata", {})}
        for m in threads.get(thread_id, [])
        if m.get("role") in ("user", "assistant")
    ]


@app.post("/api/chat")
async def chat(req: ChatRequest) -> StreamingResponse:
    thread_id = req.thread_id
    user_message = req.message

    if thread_id not in threads:
        threads[thread_id] = []

    threads[thread_id].append({"role": "user", "content": user_message})

    conversation_parts: list[str] = []
    for m in threads[thread_id]:
        if m["role"] == "user":
            conversation_parts.append(f"User: {m['content']}")
        elif m["role"] == "assistant":
            conversation_parts.append(f"Assistant: {m['content']}")

    input_text = "\n".join(conversation_parts) if len(conversation_parts) > 1 else user_message

    entity_ids = re.findall(r"pt-\d{4}", user_message)

    assert graph_client is not None  # noqa: S101

    agent = GraphAugmentedAgent(
        "triage-agent",
        client=llm_client,
        config=agent_config,
        tools=registry,
        graph_client=graph_client,
        entity_ids=entity_ids,
        decision_type="triage",
    )

    async def event_stream():  # noqa: ANN202
        full_response = ""
        run_id = ""
        graph_context_info: dict[str, Any] = {}
        tool_calls: list[dict[str, str]] = []
        try:
            async for event in agent.run_stream(input_text):
                if isinstance(event, TextDelta):
                    full_response += event.text
                    yield f"data: {json.dumps({'type': 'token', 'token': event.text})}\n\n"
                elif isinstance(event, ToolCallStart):
                    data = json.dumps({"type": "tool_start", "tool": event.tool_name, "arguments": event.arguments})
                    yield f"data: {data}\n\n"
                elif isinstance(event, ToolCallEnd):
                    tool_calls.append({"tool": event.tool_name, "result": event.result})
                    data = json.dumps({"type": "tool_end", "tool": event.tool_name, "result": event.result})
                    yield f"data: {data}\n\n"
                elif isinstance(event, AgentComplete):
                    run_id = event.metadata.get("run_id", "")
                    graph_context_info = event.metadata.get("graph_context", {})
                    meta: dict[str, Any] = {
                        "iterations": event.iterations,
                        "tool_calls_made": event.tool_calls_made,
                        "graph_context": graph_context_info,
                        "tool_calls": tool_calls,
                        "run_id": run_id,
                    }
                    threads[thread_id].append({"role": "assistant", "content": event.output, "metadata": meta})
                    yield f"data: {json.dumps({'type': 'complete', 'run_id': run_id, 'metadata': meta})}\n\n"
                elif isinstance(event, AgentError):
                    threads[thread_id].append({"role": "assistant", "content": event.output})
                    yield f"data: {json.dumps({'type': 'error', 'error': event.error})}\n\n"
        except Exception as exc:
            yield f"data: {json.dumps({'type': 'error', 'error': str(exc)})}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


@app.post("/api/feedback")
async def submit_feedback(req: FeedbackRequest) -> dict[str, str]:
    assert graph_client is not None  # noqa: S101
    score = 1.0 if req.satisfied else -1.0
    ftype = FeedbackType.THUMBS_UP if req.satisfied else FeedbackType.THUMBS_DOWN
    feedback_id = await record_feedback(graph_client, run_id=req.run_id, feedback_type=ftype, score=score)
    return {"feedback_id": feedback_id, "status": "recorded"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8007)
