app/graph.py

import os
import json
from typing import Optional, TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

from app.llm import get_llm, CombinedOutput, HPI_FIELDS, ROS_REQUIRED
from app.schemas import ClinicalBrief, HPI

_MOCK = lambda: os.environ.get("MOCK_LLM", "true").lower() == "true"


def add_messages(left: list[dict], right: list[dict]) -> list[dict]:
    return left + right


class IntakeState(TypedDict):
    messages: Annotated[list[dict], add_messages]
    clinical_state: str          # JSON of CombinedOutput (accumulated clinical data)
    missing_fields: list[str]
    current_node: str
    clinical_brief: Optional[dict]
    frontend_stage: str          # 'intake', 'hpi', 'ros', 'done'
    ros_stuck_count: int         # consecutive turns stuck in ROS with no progress


EMERGENCY_PHRASES = [
    "crushing chest pain", "can't breathe", "cannot breathe",
    "heart attack", "suicide", "kill myself", "can't move", "dying"
]


# ------------------------------------------------------------------ helpers --

def format_transcript(messages: list[dict]) -> str:
    lines = []
    for m in messages:
        role = "AI" if m["role"] == "assistant" else "Patient"
        lines.append(f"{role}: {m['content']}")
    return "\n".join(lines)


def compute_stage(state: CombinedOutput) -> str:
    if not state.chief_complaint:
        return "intake"
    for f in HPI_FIELDS:
        if not getattr(state, f):
            return "hpi"
    if len(state.ros) < ROS_REQUIRED:
        return "ros"
    return "done"


def missing_from(state: CombinedOutput) -> list[str]:
    missing = []
    if not state.chief_complaint:
        missing.append("chief complaint")
        return missing
    for f in HPI_FIELDS:
        if not getattr(state, f):
            missing.append(f"HPI:{f}")
    if len(state.ros) < ROS_REQUIRED:
        missing.append(f"ROS ({ROS_REQUIRED - len(state.ros)} more systems needed)")
    return missing


def _get_last_user_message(msgs: list[dict]) -> str:
    for m in reversed(msgs):
        if m.get("role") == "user":
            return m.get("content", "")
    return ""


def _detect_repeat(msgs: list[dict], new_reply: str) -> bool:
    """Return True if new_reply is identical to the last two stored assistant replies."""
    assistant_replies = [m.get("content", "") for m in msgs if m.get("role") == "assistant"]
    if len(assistant_replies) >= 2:
        return new_reply == assistant_replies[-1] == assistant_replies[-2]
    return False


# ------------------------------------------------------------------- nodes ---

def triage_node(state: IntakeState) -> dict:
    """Fast keyword check — no LLM call. Abort immediately on emergency phrases."""
    msgs = state.get("messages", [])
    if msgs and msgs[-1]["role"] == "user":
        content = msgs[-1]["content"].lower()
        if any(p in content for p in EMERGENCY_PHRASES):
            return {
                "messages": [{
                    "role": "assistant",
                    "content": (
                        "🚨 EMERGENCY: Your symptoms require immediate attention. "
                        "Please call 911 or go to your nearest emergency room right away."
                    )
                }],
                "current_node": "done",
                "frontend_stage": "done",
            }
    return {"current_node": "agent"}


def agent_node(state: IntakeState) -> dict:
    """
    Core agent — one LLM call per turn.
    Extracts clinical data, generates next question, builds brief when complete.
    """
    msgs = state.get("messages", [])

    # First call: no messages yet → return opening greeting
    if not msgs or (len(msgs) == 1 and msgs[0]["role"] == "assistant"):
        return {
            "messages": [{"role": "assistant", "content": "Hello, I'm conducting your pre-visit clinical intake. What brings you in today?"}],
            "clinical_state": CombinedOutput().model_dump_json(),
            "frontend_stage": "intake",
            "current_node": "agent",
            "ros_stuck_count": 0,
        }

    if msgs[-1]["role"] == "assistant":
        return {"current_node": "agent"}

    current_json = state.get("clinical_state") or CombinedOutput().model_dump_json()
    transcript = format_transcript(msgs)
    ros_stuck_count = state.get("ros_stuck_count", 0)

    try:
        pre_state = CombinedOutput.model_validate_json(current_json)
        current_stage = compute_stage(pre_state)
    except Exception:
        current_stage = "intake"

    import time
    print(f"[{time.time():.3f}] [Graph Node] Requesting LLM inference (stage={current_stage})...")

    llm = get_llm()
    result: CombinedOutput = llm.combined_call(transcript, current_json, stage=current_stage)

    # ── ROS Hallucination Guard: max 1 new ROS system per turn ──────────
    try:
        prev_ros = json.loads(current_json).get("ros") or {}
    except Exception:
        prev_ros = {}

    new_ros_keys = [k for k in result.ros if k not in prev_ros]
    if len(new_ros_keys) > 1:
        print(f"[ROSGuard] LLM added {len(new_ros_keys)} systems in one turn: {new_ros_keys}. Keeping first only.")
        allowed_ros = dict(prev_ros)
        allowed_ros[new_ros_keys[0]] = result.ros[new_ros_keys[0]]
        result = result.model_copy(update={"ros": allowed_ros})

    # ── Loop Guard ───────────────────────────────────────────────────────
    try:
        prev_state_obj = CombinedOutput.model_validate_json(current_json)
        prev_filled = sum(1 for f in HPI_FIELDS if getattr(prev_state_obj, f, None)) + len(prev_state_obj.ros)
        new_filled = sum(1 for f in HPI_FIELDS if getattr(result, f, None)) + len(result.ros)
        made_progress = new_filled > prev_filled
    except Exception:
        made_progress = True  # assume progress on parse error

    hpi_complete = all(getattr(result, f, None) for f in HPI_FIELDS)

    if not made_progress:
        last_user_msg = _get_last_user_message(msgs)

        if not hpi_complete:
            # HPI stuck — force-fill the first empty field
            for stuck_field in HPI_FIELDS:
                if not getattr(result, stuck_field, None):
                    result = result.model_copy(update={stuck_field: last_user_msg or "not specified"})
                    print(f"[LoopGuard] Force-filled HPI '{stuck_field}' = '{last_user_msg or 'not specified'}'")
                    break
        else:
            # ROS stuck — force-store the user's answer into a pending ros_asked system
            ros_stuck_count += 1
            pending = [s for s in result.ros_asked if s not in result.ros]

            if pending:
                # Store whatever the user just said as the finding for this system
                new_ros = dict(result.ros)
                new_ros[pending[0]] = [last_user_msg] if last_user_msg else ["no symptoms reported"]
                result = result.model_copy(update={"ros": new_ros})
                print(f"[LoopGuard] Force-stored ROS['{pending[0]}'] = [{last_user_msg[:40]}]")
            elif ros_stuck_count >= 2:
                # LLM isn't even updating ros_asked — force a dummy system to unblock
                stub_key = f"general_{len(result.ros)}"
                new_ros = dict(result.ros)
                new_ros[stub_key] = [last_user_msg] if last_user_msg else ["no additional symptoms"]
                result = result.model_copy(update={"ros": new_ros})
                print(f"[LoopGuard] Force-added stub ROS['{stub_key}'] after {ros_stuck_count} stuck turns.")
                ros_stuck_count = 0
    else:
        ros_stuck_count = 0  # reset counter when progress is made

    print(f"[{time.time():.3f}] [Graph Node] LLM returned. Preparing node dictionaries...")

    stage = compute_stage(result)
    missing = missing_from(result)
    reply = result.reply or "Could you tell me more?"

    # Sanitize reply — avoid storing empty or whitespace-only replies
    if not reply.strip():
        reply = "Could you tell me more?"

    # All fields complete — build the brief inline
    if stage == "done":
        from datetime import datetime, timezone
        brief = ClinicalBrief(
            chief_complaint=result.chief_complaint or "Not specified",
            hpi=HPI(
                onset=result.onset or "Not specified",
                location=result.location or "Not specified",
                duration=result.duration or "Not specified",
                character=result.character or "Not specified",
                severity=result.severity or "Not specified",
                aggravating=result.aggravating or "Not specified",
                relieving=result.relieving or "Not specified",
            ),
            ros=result.ros,
            generated_at=datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        )
        return {
            "messages": [{"role": "assistant", "content": "Your clinical summary is ready. Please wait for the doctor."}],
            "clinical_state": result.model_dump_json(),
            "missing_fields": [],
            "frontend_stage": "done",
            "current_node": "done",
            "clinical_brief": brief.model_dump(),
            "ros_stuck_count": 0,
        }

    return {
        "messages": [{"role": "assistant", "content": reply}],
        "clinical_state": result.model_dump_json(),
        "missing_fields": missing,
        "frontend_stage": stage,
        "current_node": "agent",
        "ros_stuck_count": ros_stuck_count,
    }


# -------------------------------------------------------------- graph build --

def build_graph():
    workflow = StateGraph(IntakeState)

    workflow.add_node("triage", triage_node)
    workflow.add_node("agent", agent_node)

    def route_triage(state: IntakeState) -> str:
        return state.get("current_node", "agent")

    workflow.add_edge(START, "triage")
    workflow.add_conditional_edges("triage", route_triage, {"done": END, "agent": "agent"})
    workflow.add_edge("agent", END)

    checkpointer = MemorySaver()
    graph = workflow.compile(checkpointer=checkpointer)
    # NOTE: interrupt_after removed — state accumulates via MemorySaver reducer
    # on every fresh invoke, which is correct behavior (has_next is always False)

    return graph, checkpointer