DentalManagementMH07/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>(),
});

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)\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 END;
}

// ── 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! Would you like to reschedule?",
  Spanish:          "¡Lo entendemos! ¿Le gustaría reprogramar su cita?",
  Portuguese:       "Entendemos! Gostaria de reagendar a sua consulta?",
  Mandarin:         "我们理解！您想重新安排预约吗？",
  Cantonese:        "我們理解！您想重新安排預約嗎？",
  Arabic:           "نتفهم ذلك! هل تود إعادة جدولة موعدك؟",
  "Haitian Creole": "Nou konprann! Èske ou ta renmen repwograme randevou ou?",
};

// ── 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} regardless of the language the patient used. 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 };
  }
}

async function rescheduleNode(state: GraphStateType, config: any) {
  const apiKey: string | undefined = config?.configurable?.apiKey;
  const lang    = state.language || "English";
  const fallback = RESCHEDULE_FALLBACKS[lang] ?? RESCHEDULE_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 assistant. The patient cannot make their appointment. Write a short, empathetic SMS reply (1 sentence max) that says it is understandable and asks if they would like to reschedule. You MUST reply in ${lang} regardless of the language the patient used. 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 };
  }
}

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

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

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