DentalManagementMH05/apps/Backend/src/ai/reminder-graph.ts

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

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

type GraphStateType = typeof GraphState.State;

// ── Intent classifier — multilingual yes/no keywords ─────────────────────────

function classifyNode(state: GraphStateType) {
  const text = state.message.toLowerCase().trim();

  const yesPatterns = /\b(yes|yeah|yep|yup|sure|ok|okay|confirmed|confirm|will be there|sounds good|see you|great|perfect|absolutely|definitely|sí|si|claro|por supuesto|confirmo|de acuerdo|seguro|estaré|sim|confirmado|com certeza|好的|确认|可以|好|明白|نعم|حسنا|موافق|wi|dakò|oke)\b/;
  const noPatterns  = /\b(no|nope|can't|cannot|won't|not available|unavailable|cancel|reschedule|busy|sorry|unable|not coming|not going|no puedo|no podré|cancelar|reprogramar|ocupado|lo siento|não posso|não vou|reagendar|占线|无法|取消|لا|لا أستطيع|إلغاء|pa kapab|pa ka|can not make|cannot make|not make|make it)\b/;

  if (yesPatterns.test(text)) return { intent: "yes" };
  if (noPatterns.test(text))  return { intent: "no" };
  return { intent: "other" };
}

function routeByIntent(state: GraphStateType): string {
  if (state.intent === "yes") return "confirm";
  if (state.intent === "no")  return "reschedule";
  return "other";
}

// ── Confirmation fallbacks (with appointment datetime) ────────────────────────

function buildConfirmFallback(lang: string, apptDatetime: string): string {
  const appt = apptDatetime || "your scheduled appointment";
  const map: Record<string, string> = {
    English:          `Thank you for your confirmation! We look forward to seeing you on ${appt}.`,
    Spanish:          `¡Gracias por confirmar! Le esperamos el ${appt}.`,
    Portuguese:       `Obrigado por confirmar! Aguardamos a sua visita em ${appt}.`,
    Mandarin:         `感谢您的确认！我们期待在 ${appt} 见到您。`,
    Cantonese:        `感謝您的確認！我們期待在 ${appt} 見到您。`,
    Arabic:           `شكراً على تأكيدك! نتطلع إلى رؤيتك في ${appt}.`,
    "Haitian Creole": `Mèsi dèske ou konfime! N'ap tann ou ${appt}.`,
  };
  return map[lang] ?? map["English"]!;
}

// ── Reschedule fallbacks ──────────────────────────────────────────────────────

const RESCHEDULE_FALLBACKS: Record<string, string> = {
  English:          "It is understandable! When would you like to reschedule?",
  Spanish:          "¡Lo entendemos! ¿Cuándo le gustaría reprogramar su cita?",
  Portuguese:       "Entendemos! Quando gostaria de reagendar a sua consulta?",
  Mandarin:         "我们理解！您想什么时候重新安排预约？",
  Cantonese:        "我們理解！您想幾時重新安排預約？",
  Arabic:           "نتفهم ذلك! متى تود إعادة جدولة موعدك؟",
  "Haitian Creole": "Nou konprann! Ki lè ou ta renmen repwograme randevou ou?",
};

// ── New-appointment fallbacks (other intent, appointment keywords detected) ───

const NEW_APPT_FALLBACKS: Record<string, string> = {
  English:          "Of course! Are you a new patient or an existing patient?",
  Spanish:          "¡Por supuesto! ¿Es usted un paciente nuevo o ya ha visitado nuestra clínica?",
  Portuguese:       "Claro! Você é um paciente novo ou já veio ao nosso consultório antes?",
  Mandarin:         "当然！您是新患者还是我们的现有患者？",
  Cantonese:        "當然！您係新病人定係我們的舊病人？",
  Arabic:           "بالطبع! هل أنت مريض جديد أم مريض حالي؟",
  "Haitian Creole": "Byensèten! Èske ou se yon nouvo pasyan oswa yon pasyan egzistan?",
};

// ── General-other fallbacks ───────────────────────────────────────────────────

const GENERAL_FALLBACKS: Record<string, string> = {
  English:          "Thank you for your message! Our office staff will be happy to assist you shortly.",
  Spanish:          "¡Gracias por su mensaje! El personal de nuestra oficina estará encantado de ayudarle en breve.",
  Portuguese:       "Obrigado pela sua mensagem! Nossa equipe terá prazer em ajudá-lo em breve.",
  Mandarin:         "感谢您的留言！我们的办公室工作人员将很快为您提供帮助。",
  Cantonese:        "感謝您的留言！我們的辦公室工作人員將很快為您提供幫助。",
  Arabic:           "شكراً على رسالتك! سيسعد فريق مكتبنا بمساعدتك قريباً.",
  "Haitian Creole": "Mèsi pou mesaj ou! Ekip biwo nou an pral kontan ede ou byento.",
};

// ── LangGraph nodes ───────────────────────────────────────────────────────────

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

  if (!apiKey) return { reply: fallback };

  try {
    const llm = new ChatGoogleGenerativeAI({ model: "gemini-1.5-flash", apiKey });
    const apptClause = appt ? ` Their appointment is on ${appt}.` : "";
    const response = await llm.invoke([
      {
        role: "system",
        content:
          `You are a friendly dental office assistant. Write a short, warm SMS reply (1-2 sentences max) thanking the patient for confirming their appointment and reminding them of the date and time.${apptClause} You MUST reply in ${lang}. Do not add any formatting or extra text.`,
      },
      { role: "user", content: `Patient replied: "${state.message}"` },
    ]);
    return { reply: String(response.content) || fallback };
  } catch {
    return { reply: fallback };
  }
}

/**
 * Strip any leading self-introduction ("Hi! My name is Lisa...") from a
 * reschedule greeting template so it can be used as MSG 2 without duplicating
 * the intro that was already sent as MSG 1.
 */
function stripIntroFromFallback(text: string): string {
  // Matches patterns like: "Hi! My name is Lisa, ... . " or "Hi, I'm Lisa ... . "
  const stripped = text
    .replace(/^(Hi[!,]?\s*)?(My name is|I'?m|I am)\s+[^.!?]+[.!?]\s*/i, "")
    .trim();
  return stripped || text;
}

async function rescheduleNode(state: GraphStateType, config: any) {
  const apiKey: string | undefined = config?.configurable?.apiKey;
  const lang     = state.language || "English";
  // Strip any self-intro from the saved template — intro is already sent as MSG 1
  const rawFallback = state.rescheduleGreeting || (RESCHEDULE_FALLBACKS[lang] ?? RESCHEDULE_FALLBACKS["English"]!);
  const fallback    = stripIntroFromFallback(rawFallback);

  if (!apiKey) return { reply: fallback };

  try {
    const llm = new ChatGoogleGenerativeAI({ model: "gemini-1.5-flash", apiKey });
    const response = await llm.invoke([
      {
        role: "system",
        content:
          `You are a friendly dental office assistant. The patient cannot make their appointment. Write a short, empathetic SMS reply (1-2 sentences max) acknowledging they can't make it and asking when they would like to reschedule. Do NOT introduce yourself or say your name — the introduction has already been sent in a separate message. You MUST reply in ${lang}. No formatting, no extra text.`,
      },
      { role: "user", content: `Patient replied: "${state.message}"` },
    ]);
    return { reply: String(response.content) || fallback };
  } catch {
    return { reply: fallback };
  }
}

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

  const isAppointmentRequest = /appointment|schedule|book|come in|visit|check.?up|cleaning|tooth|teeth|pain|dental|another|new appt/i.test(text);

  if (isAppointmentRequest) {
    const fallback = NEW_APPT_FALLBACKS[lang] ?? NEW_APPT_FALLBACKS["English"]!;
    if (!apiKey) return { reply: fallback, intent: "wants_appointment" };
    try {
      const llm = new ChatGoogleGenerativeAI({ model: "gemini-1.5-flash", apiKey });
      const response = await llm.invoke([
        {
          role: "system",
          content: `You are a friendly dental office AI assistant. The patient wants to schedule an appointment. Ask them in ${lang} whether they are a new patient or an existing patient. One sentence, no formatting.`,
        },
        { role: "user", content: `Patient said: "${state.message}"` },
      ]);
      return { reply: String(response.content) || fallback, intent: "wants_appointment" };
    } catch {
      return { reply: fallback, intent: "wants_appointment" };
    }
  }

  const fallback = state.generalFallback || (GENERAL_FALLBACKS[lang] ?? GENERAL_FALLBACKS["English"]!);
  if (!apiKey) return { reply: fallback };
  try {
    const llm = new ChatGoogleGenerativeAI({ model: "gemini-1.5-flash", apiKey });
    const response = await llm.invoke([
      {
        role: "system",
        content: `You are a friendly dental office AI assistant. Respond helpfully to the patient's message in ${lang}. Keep it to 1-2 sentences, no formatting. For non-dental questions, let them know our office staff can assist.`,
      },
      { role: "user", content: `Patient said: "${state.message}"` },
    ]);
    return { reply: String(response.content) || fallback };
  } catch {
    return { reply: fallback };
  }
}

// ── Graph ─────────────────────────────────────────────────────────────────────

const graph = new StateGraph(GraphState)
  .addNode("classify",   classifyNode)
  .addNode("confirm",    confirmNode)
  .addNode("reschedule", rescheduleNode)
  .addNode("other",      otherNode)
  .addEdge(START, "classify")
  .addConditionalEdges("classify", routeByIntent, {
    confirm:    "confirm",
    reschedule: "reschedule",
    other:      "other",
  })
  .addEdge("confirm",    END)
  .addEdge("reschedule", END)
  .addEdge("other",      END)
  .compile();

export async function runReminderGraph(
  patientMessage:      string,
  apiKey:              string,
  language            = "English",
  appointmentDatetime = "",
  rescheduleGreeting  = "",
  generalFallback     = ""
): Promise<{ reply: string | null; intent: string | null }> {
  const result = await graph.invoke(
    { message: patientMessage, intent: "", reply: "", language, appointmentDatetime, rescheduleGreeting, generalFallback },
    { configurable: { apiKey } }
  );
  return {
    reply:  result.reply  || null,
    intent: result.intent || null,
  };
}