DentalManagement09/apps/Backend/src/ai/internal-chat-graph.ts

import { ChatGoogleGenerativeAI } from "@langchain/google-genai";

// ─── Intent types ─────────────────────────────────────────────────────────────

export type InternalChatIntent =
  | "check_eligibility"     // by patient name → look up in DB
  | "eligibility_by_id"     // by explicit memberId + dob (no name)
  | "check_and_claim"       // eligibility + claim procedures
  | "find_patient"          // look up patient record only
  | "schedule_appointment"  // add patient to today's (or specified) schedule
  | "claim_only"            // submit claim for procedures (no eligibility check)
  | "navigate_claims"
  | "navigate_schedule"
  | "general";

export interface ChatClassification {
  intent: InternalChatIntent;
  // --- patient resolution (one of name OR id+dob) ---
  patientName?: string;       // for check_eligibility / find_patient / schedule_appointment
  memberId?: string;          // for eligibility_by_id / check_and_claim
  dob?: string;               // for eligibility_by_id / check_and_claim (MM/DD/YYYY)
  // --- insurance hint (only if explicitly stated in the message) ---
  insuranceHint?: string;     // raw text, e.g. "masshealth", "BCBS", "CCA"
  // --- procedures (raw text, NOT CDT codes — CDT lookup is done in workflow) ---
  procedureNames?: string[];  // for check_and_claim, e.g. ["perio exam", "adult cleaning"]
  // --- scheduling ---
  appointmentDate?: string;   // for schedule_appointment, YYYY-MM-DD (omit = today)
  appointmentTime?: string;   // for schedule_appointment, HH:MM 24h (omit = 09:00)
  fallbackReply: string;
}

// ─── System prompt ────────────────────────────────────────────────────────────

const BASE_SYSTEM_PROMPT = `You are an internal assistant for a dental office management app.
Staff type natural language commands. Your ONLY job is to classify the intent and extract
structured parameters. Do NOT map procedure names to CDT codes — return them as plain text.

Respond ONLY with valid JSON (no markdown fences):
{
  "intent": "<intent>",
  "patientName": "<full name if mentioned by name>",
  "memberId": "<member/insurance ID if given explicitly>",
  "dob": "<date of birth in MM/DD/YYYY if given>",
  "insuranceHint": "<insurance name only if explicitly stated in the message, e.g. 'masshealth', 'BCBS MA', 'CCA'>",
  "procedureNames": ["<raw procedure name>", ...],
  "appointmentDate": "<YYYY-MM-DD if a specific date is mentioned, omit for today>",
  "appointmentTime": "<HH:MM 24h if a specific time is mentioned, omit if not stated>",
  "fallbackReply": "<1-2 sentence reply to show the user>"
}

Omit any field that is not present in the message.

Intents:
- check_eligibility     : user wants to check insurance for a patient identified by NAME only
                          e.g. "check Maria Jesus", "verify insurance for John Smith"
- eligibility_by_id     : user provides a member ID and date of birth (no patient name)
                          e.g. "check masshealth for 100xxxx, 10/10/1988"
- check_and_claim       : user wants to check eligibility AND submit procedures as claims
                          e.g. "check masshealth for 100xxxx, 10/10/1988 and claim perio exam and adult cleaning"
                          e.g. "check Maria Jesus and claim D0120 D1110"
- find_patient          : look up a patient record only, no eligibility
                          e.g. "find patient John", "look up Smith"
- schedule_appointment  : add a patient to the schedule (today or a specified date/time)
                          e.g. "put John Smith in today's schedule"
                          e.g. "schedule Maria at 2pm tomorrow"
                          e.g. "add Jane Doe at 10:30"
- claim_only            : submit a claim for procedures WITHOUT an eligibility check
                          e.g. "claim comprehensive exam and Pano for her"
                          e.g. "claim D0120 and D1110 for John Smith"
                          e.g. "bill adult cleaning for Maria"
                          Use this when no eligibility check is requested — just billing/claiming services
- navigate_claims       : open the claims page
- navigate_schedule     : open the appointments/schedule page
- general               : anything else

Rules:
- For check_and_claim and claim_only, procedureNames should be the RAW user text
  (e.g. "perio exam", "adult cleaning", "D0120") — do NOT translate to codes
- insuranceHint is only set when the user explicitly names an insurance in the message
- Keep fallbackReply to 1-2 sentences
- For navigate intents, fallbackReply = "Opening the [page] page..."
- For schedule_appointment, appointmentDate omitted means today; appointmentTime omitted means no preference
- IMPORTANT: Use the conversation history to resolve pronouns and references.
  If the user says "her", "him", "them", "the patient", or "same patient", look back through
  the conversation history to find the patient name that was mentioned most recently.
  Always populate patientName (or memberId) from history when a pronoun is used.
  Never return an empty patientName just because the current message uses a pronoun.`;

// ─── Classifier ───────────────────────────────────────────────────────────────

export async function classifyInternalChat(
  message: string,
  apiKey: string,
  extraSystemPrompt?: string,
  history: { role: "user" | "assistant"; text: string }[] = []
): Promise<ChatClassification> {
  const fallback: ChatClassification = {
    intent: "general",
    fallbackReply:
      "I can search for a patient, check eligibility, run check & claim, schedule appointments, or navigate to claims or appointments.",
  };

  if (!apiKey) return fallback;

  const systemPrompt = extraSystemPrompt
    ? `${BASE_SYSTEM_PROMPT}\n\nAdditional office context:\n${extraSystemPrompt}`
    : BASE_SYSTEM_PROMPT;

  try {
    const llm = new ChatGoogleGenerativeAI({ model: "gemini-flash-latest", apiKey });

    // Gemini requires conversation to start with a user turn — drop any leading assistant messages
    const trimmedHistory = history.slice(
      history.findIndex((h) => h.role === "user")
    ).filter((_, i, arr) => {
      // Also drop consecutive same-role messages (keep last of each run)
      if (i === arr.length - 1) return true;
      return arr[i]!.role !== arr[i + 1]!.role;
    });

    const historyMessages = trimmedHistory.map((h) => ({
      role: h.role,
      content: h.text,
    }));
    const response = await llm.invoke([
      { role: "system", content: systemPrompt },
      ...historyMessages,
      { role: "user", content: message },
    ]);

    const raw = String(response.content).trim();
    const jsonStr = raw.replace(/^```json\s*/i, "").replace(/```\s*$/, "").trim();
    const parsed = JSON.parse(jsonStr) as ChatClassification;
    if (!parsed.intent || !parsed.fallbackReply) return fallback;
    return parsed;
  } catch {
    return fallback;
  }
}