DentalManagementMH05/apps/Backend/src/ai/new-patient-graph.ts

import { StateGraph, END, START, Annotation } from "@langchain/langgraph";
import { ChatGoogleGenerativeAI } from "@langchain/google-genai";
import type { ConversationStage } from "./aiHandoffStore";

// ── Graph state ───────────────────────────────────────────────────────────────

const GraphState = Annotation.Root({
  message:   Annotation<string>(),
  stage:     Annotation<string>(),
  intent:    Annotation<string>(),
  reply:     Annotation<string>(),
  language:  Annotation<string>(),
  nextStage: Annotation<string>(),
});

type GraphStateType = typeof GraphState.State;

// ── Transfer-to-staff fallback (multilingual) ─────────────────────────────────

const TRANSFER: Record<string, string> = {
  English:          "Thank you for reaching out! Our office staff will assist you shortly.",
  Spanish:          "¡Gracias por comunicarse! Un miembro del personal de la oficina le atenderá en breve.",
  Portuguese:       "Obrigado pelo contato! Nossa equipe entrará em contato em breve.",
  Mandarin:         "感谢您的联系！我们的工作人员将很快为您提供帮助。",
  Cantonese:        "感謝您的聯繫！我們的工作人員將很快為您提供幫助。",
  Arabic:           "شكراً لتواصلك! سيتواصل معك أحد موظفي المكتب قريباً.",
  "Haitian Creole": "Mèsi dèske ou kontakte nou! Yon manm ekip biwo a pral ede ou byento.",
};

function transferMsg(lang: string): string {
  return TRANSFER[lang] ?? TRANSFER["English"]!;
}

// ── Intent classifiers ────────────────────────────────────────────────────────

function wantsAppointment(text: string): boolean {
  return /appointment|schedule|cleaning|checkup|teeth|tooth|cavity|pain|dental|filling|crown|whitening|x-ray|exam|hygienist|dentist/i.test(text);
}

function isNewPatient(text: string): boolean {
  return /new patient|first time|first visit|never been|brand new|haven't been|i am new/i.test(text);
}

function isExistingPatient(text: string): boolean {
  return /existing|been there|have been|already|before|i have been|returning|came before|i was there/i.test(text);
}

function hasMassHealth(text: string): boolean {
  return /masshealth|mass health|medicaid|masscare/i.test(text);
}

function hasOtherInsurance(text: string): boolean {
  return /blue cross|delta dental|cigna|aetna|united|metlife|guardian|humana|tufts|harvard pilgrim|bmchp|yes|i have|my insurance|i do/i.test(text);
}

function hasNoInsurance(text: string): boolean {
  return /no insurance|uninsured|self.pay|self pay|i don't|don't have|no i don't|i have no/i.test(text);
}

function sameInsurance(text: string): boolean {
  return /yes|same|still have|haven't changed|no change/i.test(text);
}

function changedInsurance(text: string): boolean {
  return /no|changed|different|new insurance|switched|lost|expired/i.test(text);
}

// ── LLM reply helper ──────────────────────────────────────────────────────────

async function llmReply(
  system: string,
  userMsg: string,
  fallback: string,
  apiKey: string
): Promise<string> {
  try {
    const llm = new ChatGoogleGenerativeAI({ model: "gemini-1.5-flash", apiKey });
    const res = await llm.invoke([
      { role: "system", content: system },
      { role: "user",   content: userMsg },
    ]);
    return String(res.content).trim() || fallback;
  } catch {
    return fallback;
  }
}

// ── Graph nodes ───────────────────────────────────────────────────────────────

// Classify intent based on current stage
function classifyNode(state: GraphStateType) {
  const text  = state.message.toLowerCase();
  const stage = state.stage as ConversationStage;

  let intent = "other";

  if (stage === "new_patient_greeted") {
    intent = wantsAppointment(text) ? "wants_appointment" : "other";
  } else if (stage === "asked_new_or_existing") {
    if (isNewPatient(text))      intent = "new_patient";
    else if (isExistingPatient(text)) intent = "existing_patient";
  } else if (stage === "asked_new_patient_insurance") {
    if (hasMassHealth(text))     intent = "masshealth";
    else if (hasNoInsurance(text)) intent = "no_insurance";
    else if (hasOtherInsurance(text)) intent = "other_insurance";
  } else if (stage === "asked_existing_insurance") {
    if (sameInsurance(text))     intent = "same_insurance";
    else if (changedInsurance(text)) intent = "changed_insurance";
  } else if (stage === "asked_appointment_time") {
    intent = "appointment_time";
  } else if (stage === "asked_appointment_preference") {
    intent = "appointment_preference_reply";
  } else if (stage === "asked_self_pay") {
    intent = "self_pay_reply";
  }

  return { intent };
}

function routeNode(state: GraphStateType): string {
  const stage  = state.stage as ConversationStage;
  const intent = state.intent;

  if (stage === "new_patient_greeted")         return intent === "wants_appointment" ? "ask_new_or_existing" : "transfer";
  if (stage === "asked_new_or_existing")       return intent === "new_patient" ? "ask_new_patient_insurance" : intent === "existing_patient" ? "ask_existing_insurance" : "transfer";
  if (stage === "asked_new_patient_insurance") return intent === "masshealth" ? "ask_masshealth_info" : intent === "no_insurance" ? "ask_appointment_time" : "transfer";
  if (stage === "asked_existing_insurance")    return intent === "same_insurance" ? "ask_appointment_time" : "transfer";
  if (stage === "asked_appointment_time")          return "acknowledge_appointment_time";
  if (stage === "asked_appointment_preference")    return "handle_appointment_preference";
  if (stage === "asked_self_pay")                  return "handle_self_pay";

  return "transfer";
}

// ── Response nodes ────────────────────────────────────────────────────────────

async function askNewOrExistingNode(state: GraphStateType, config: any) {
  const lang    = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "Are you a new patient or an existing patient?",
    Spanish:          "¿Es usted un paciente nuevo o ya ha visitado nuestra clínica antes?",
    Portuguese:       "Você é um paciente novo ou já veio ao nosso consultório antes?",
    Mandarin:         "您是新患者还是现有患者？",
    Cantonese:        "您係新病人定係舊病人？",
    Arabic:           "هل أنت مريض جديد أم مريض حالي؟",
    "Haitian Creole": "Èske ou se yon nouvo pasyan oswa yon pasyan egzistan?",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. Ask the patient in ${lang} whether they are a new patient or an existing patient. One sentence, no formatting.`,
        `Patient wants an appointment. Ask if new or existing.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "asked_new_or_existing" };
}

async function askNewPatientInsuranceNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "Do you have any dental insurance?",
    Spanish:          "¿Tiene seguro dental?",
    Portuguese:       "Você tem plano odontológico?",
    Mandarin:         "您有牙科保险吗？",
    Cantonese:        "您有牙科保險嗎？",
    Arabic:           "هل لديك تأمين أسنان؟",
    "Haitian Creole": "Èske ou gen asirans dantè?",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. Ask the new patient in ${lang} if they have dental insurance. One sentence, no formatting.`,
        `New patient confirmed. Ask about insurance.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "asked_new_patient_insurance" };
}

async function askExistingInsuranceNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "Do you still have the same insurance?",
    Spanish:          "¿Sigue teniendo el mismo seguro?",
    Portuguese:       "Você ainda tem o mesmo plano?",
    Mandarin:         "您还有相同的保险吗？",
    Cantonese:        "您仍然有相同的保險嗎？",
    Arabic:           "هل لا تزال تمتلك نفس التأمين؟",
    "Haitian Creole": "Èske ou toujou gen menm asirans lan?",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. Ask the existing patient in ${lang} if they still have the same dental insurance on file. One sentence, no formatting.`,
        `Existing patient confirmed. Ask about insurance.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "asked_existing_insurance" };
}

async function askMassHealthInfoNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "I can check your MassHealth coverage! Please text me your Member ID and date of birth.",
    Spanish:          "¡Puedo verificar su cobertura de MassHealth! Por favor envíeme su número de miembro y fecha de nacimiento.",
    Portuguese:       "Posso verificar sua cobertura MassHealth! Por favor envie seu número de membro e data de nascimento.",
    Mandarin:         "我可以查看您的MassHealth保险！请发送您的会员ID和出生日期。",
    Cantonese:        "我可以查核您的MassHealth保險！請傳送您的會員ID和出生日期。",
    Arabic:           "يمكنني التحقق من تغطية MassHealth الخاصة بك! من فضلك أرسل لي رقم العضوية وتاريخ الميلاد.",
    "Haitian Creole": "Mwen ka verifye asirans MassHealth ou! Tanpri voye ID manm ou ak dat nesans ou.",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. The patient has MassHealth. Tell them in ${lang} that you can check their coverage and ask them to send their Member ID and date of birth. 1-2 sentences, no formatting.`,
        `Patient has MassHealth. Ask for member ID and DOB.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "awaiting_masshealth_info" };
}

async function askAppointmentTimeNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "When would you like to make an appointment?",
    Spanish:          "¿Cuándo le gustaría hacer una cita?",
    Portuguese:       "Quando você gostaria de agendar uma consulta?",
    Mandarin:         "您想什么时候预约？",
    Cantonese:        "您想幾時預約？",
    Arabic:           "متى تريد تحديد موعد؟",
    "Haitian Creole": "Ki lè ou ta renmen fè yon randevou?",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. Ask the patient in ${lang} when they would like to schedule their appointment. One sentence, no formatting.`,
        `Ask when to schedule.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "asked_appointment_time" };
}

async function acknowledgeAppointmentTimeNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "Thank you! Our office staff will confirm your appointment details shortly.",
    Spanish:          "¡Gracias! El personal de la oficina confirmará los detalles de su cita en breve.",
    Portuguese:       "Obrigado! Nossa equipe confirmará os detalhes da sua consulta em breve.",
    Mandarin:         "谢谢！我们的工作人员将很快确认您的预约详情。",
    Cantonese:        "多謝！我們的工作人員將很快確認您的預約詳情。",
    Arabic:           "شكراً! سيؤكد فريق مكتبنا تفاصيل موعدك قريباً.",
    "Haitian Creole": "Mèsi! Ekip biwo nou an pral konfime detay randevou ou a byento.",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. The patient stated their preferred appointment time. Acknowledge in ${lang} and tell them the office staff will confirm shortly. 1-2 sentences, no formatting.`,
        `Patient said: "${state.message}". Acknowledge.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "done" };
}

// ── Post-Selenium: appointment preference (MassHealth ACTIVE) ─────────────────

async function handleAppointmentPreferenceNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;

  const fallbacks: Record<string, string> = {
    English:          "When would you like to come in? Are you looking for a routine check-up and teeth cleaning, or do you have a tooth problem or pain?",
    Spanish:          "¿Cuándo le gustaría venir? ¿Busca una revisión rutinaria y limpieza dental, o tiene algún problema dental o dolor?",
    Portuguese:       "Quando gostaria de vir? Você busca uma consulta de rotina e limpeza, ou tem algum problema dentário ou dor?",
    Mandarin:         "您想什么时候来？您是想做常规检查和洗牙，还是您有牙齿问题或疼痛？",
    Cantonese:        "您想幾時來？您是想做例行檢查和洗牙，還是您有牙齒問題或疼痛？",
    Arabic:           "متى تودّ الحضور؟ هل تبحث عن فحص روتيني وتنظيف أسنان، أم أن لديك مشكلة في الأسنان أو ألماً؟",
    "Haitian Creole": "Ki lè ou ta renmen vini? Èske ou ap chèche yon egzamen woutin ak netwayaj dan, oswa èske ou gen yon pwoblèm dan oswa doulè?",
  };
  const fallback = fallbacks[lang] ?? fallbacks["English"]!;

  const reply = apiKey
    ? await llmReply(
        `You are a friendly dental office AI assistant. The patient's MassHealth is active. In ${lang}, ask when they would like to come in and whether they want a routine check-up and teeth cleaning, or if they have a dental problem or pain. 1-2 sentences, no formatting.`,
        `MassHealth is active. Ask appointment preference.`,
        fallback, apiKey
      )
    : fallback;

  return { reply, nextStage: "asked_appointment_time" };
}

// ── Post-Selenium: self-pay offer (MassHealth INACTIVE) ───────────────────────

async function handleSelfPayNode(state: GraphStateType, config: any) {
  const lang   = state.language || "English";
  const apiKey: string | undefined = config?.configurable?.apiKey;
  const text   = state.message.toLowerCase();

  // Classify yes/no for self-pay
  const acceptsSelfPay = /yes|sure|ok|okay|yep|yeah|sí|si|claro|sim|confirmado|好的|نعم|wi|oke/i.test(text);
  const declinesSelfPay = /no|nope|can't|won't|not interested|no puedo|não|لا|pa ka/i.test(text);

  if (declinesSelfPay) {
    const declineFallbacks: Record<string, string> = {
      English:          "No problem! If you ever change your mind or have any questions, feel free to reach out. Our staff is happy to help.",
      Spanish:          "¡Sin problema! Si cambia de opinión o tiene alguna pregunta, no dude en contactarnos.",
      Portuguese:       "Sem problema! Se mudar de ideia ou tiver dúvidas, entre em contato. Nossa equipe está feliz em ajudar.",
      Mandarin:         "没关系！如果您改变主意或有任何问题，随时联系我们，我们的工作人员很乐意帮助您。",
      Cantonese:        "沒問題！如果您改變主意或有任何問題，隨時聯繫我們。",
      Arabic:           "لا بأس! إذا غيّرت رأيك أو كان لديك أي سؤال، لا تتردد في التواصل معنا.",
      "Haitian Creole": "Pa gen pwoblèm! Si ou chanje lide ou oswa gen kesyon, pa ezite kontakte nou.",
    };
    const fallback = declineFallbacks[lang] ?? declineFallbacks["English"]!;
    const reply = apiKey
      ? await llmReply(
          `You are a friendly dental office AI assistant. The patient declined a self-pay appointment. In ${lang}, politely close the conversation and invite them to reach out anytime. 1-2 sentences, no formatting.`,
          `Patient declined self-pay.`,
          fallback, apiKey
        )
      : fallback;
    return { reply, nextStage: "done" };
  }

  if (acceptsSelfPay) {
    const acceptFallbacks: Record<string, string> = {
      English:          "When would you like to come in for your examination?",
      Spanish:          "¿Cuándo le gustaría venir para su examen?",
      Portuguese:       "Quando você gostaria de vir para seu exame?",
      Mandarin:         "您想什么时候来进行检查？",
      Cantonese:        "您想幾時來進行檢查？",
      Arabic:           "متى تودّ الحضور للفحص؟",
      "Haitian Creole": "Ki lè ou ta renmen vini pou egzamen ou?",
    };
    const fallback = acceptFallbacks[lang] ?? acceptFallbacks["English"]!;
    const reply = apiKey
      ? await llmReply(
          `You are a friendly dental office AI assistant. The patient agreed to a self-pay examination. In ${lang}, ask when they would like to come in. One sentence, no formatting.`,
          `Patient agreed to self-pay. Ask when to come in.`,
          fallback, apiKey
        )
      : fallback;
    return { reply, nextStage: "asked_appointment_time" };
  }

  // Ambiguous — transfer to staff
  return { reply: transferMsg(lang), nextStage: "done" };
}

function transferNode(state: GraphStateType) {
  const lang = state.language || "English";
  return { reply: transferMsg(lang), nextStage: "done" };
}

// ── Graph assembly ────────────────────────────────────────────────────────────

const graph = new StateGraph(GraphState)
  .addNode("classify",                     classifyNode)
  .addNode("ask_new_or_existing",          askNewOrExistingNode)
  .addNode("ask_new_patient_insurance",    askNewPatientInsuranceNode)
  .addNode("ask_existing_insurance",       askExistingInsuranceNode)
  .addNode("ask_masshealth_info",          askMassHealthInfoNode)
  .addNode("ask_appointment_time",         askAppointmentTimeNode)
  .addNode("acknowledge_appointment_time", acknowledgeAppointmentTimeNode)
  .addNode("handle_appointment_preference",handleAppointmentPreferenceNode)
  .addNode("handle_self_pay",              handleSelfPayNode)
  .addNode("transfer",                     transferNode)
  .addEdge(START, "classify")
  .addConditionalEdges("classify", routeNode, {
    ask_new_or_existing:           "ask_new_or_existing",
    ask_new_patient_insurance:     "ask_new_patient_insurance",
    ask_existing_insurance:        "ask_existing_insurance",
    ask_masshealth_info:           "ask_masshealth_info",
    ask_appointment_time:          "ask_appointment_time",
    acknowledge_appointment_time:  "acknowledge_appointment_time",
    handle_appointment_preference: "handle_appointment_preference",
    handle_self_pay:               "handle_self_pay",
    transfer:                      "transfer",
  })
  .addEdge("ask_new_or_existing",           END)
  .addEdge("ask_new_patient_insurance",     END)
  .addEdge("ask_existing_insurance",        END)
  .addEdge("ask_masshealth_info",           END)
  .addEdge("ask_appointment_time",          END)
  .addEdge("acknowledge_appointment_time",  END)
  .addEdge("handle_appointment_preference", END)
  .addEdge("handle_self_pay",               END)
  .addEdge("transfer",                      END)
  .compile();

// ── Public API ────────────────────────────────────────────────────────────────

export async function runNewPatientStep(
  message:  string,
  stage:    ConversationStage,
  language: string,
  apiKey:   string
): Promise<{ reply: string; nextStage: ConversationStage }> {
  const result = await graph.invoke(
    { message, stage, intent: "", reply: "", language, nextStage: "" },
    { configurable: { apiKey } }
  );
  return {
    reply:     result.reply     || transferMsg(language),
    nextStage: (result.nextStage as ConversationStage) || "done",
  };
}