helix-engage-server/src/widget/widget-chat.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { streamText, tool, stepCountIs } from 'ai';
import { z } from 'zod';
import type { LanguageModel, ModelMessage } from 'ai';
import { createAiModel, isAiConfigured, type AiProviderOpts } from '../ai/ai-provider';
import { PlatformGraphqlService } from '../platform/platform-graphql.service';
import { WidgetService } from './widget.service';
import { AiConfigService } from '../config/ai-config.service';

@Injectable()
export class WidgetChatService {
    private readonly logger = new Logger(WidgetChatService.name);
    private readonly apiKey: string;
    private knowledgeBase: string | null = null;
    private kbLoadedAt = 0;
    private readonly kbTtlMs = 5 * 60 * 1000;

    constructor(
        private config: ConfigService,
        private platform: PlatformGraphqlService,
        private widget: WidgetService,
        private aiConfig: AiConfigService,
    ) {
        this.apiKey = config.get<string>('platform.apiKey') ?? '';
        if (!this.hasAiModel()) {
            this.logger.warn('AI not configured — widget chat will return fallback replies');
        }
    }

    // Build the model on demand so admin updates to provider/model take effect
    // immediately. Construction is cheap (just wraps the SDK clients).
    private aiOpts(): AiProviderOpts {
        const cfg = this.aiConfig.getConfig();
        return {
            provider: cfg.provider,
            model: cfg.model,
            anthropicApiKey: this.config.get<string>('ai.anthropicApiKey'),
            openaiApiKey: this.config.get<string>('ai.openaiApiKey'),
        };
    }

    private buildAiModel(): LanguageModel | null {
        return createAiModel(this.aiOpts());
    }

    private get auth() {
        return `Bearer ${this.apiKey}`;
    }

    hasAiModel(): boolean {
        return isAiConfigured(this.aiOpts());
    }

    // Find-or-create a lead by phone. Delegates to WidgetService so there's
    // a single source of truth for the dedup window + lead shape across
    // chat + book + contact.
    async findOrCreateLead(name: string, phone: string): Promise<string> {
        return this.widget.findOrCreateLeadByPhone(name, phone, {
            source: 'WEBSITE',
            status: 'NEW',
            interestedService: 'Website Chat',
        });
    }

    // Fetch the first name of the lead's primary contact so we can greet the
    // visitor by name in the system prompt. Returns 'there' on any failure.
    async getLeadFirstName(leadId: string): Promise<string> {
        try {
            const data = await this.platform.queryWithAuth<any>(
                `query($id: UUID!) {
                    leads(filter: { id: { eq: $id } }, first: 1) {
                        edges { node { id contactName { firstName } } }
                    }
                }`,
                { id: leadId },
                this.auth,
            );
            const firstName = data?.leads?.edges?.[0]?.node?.contactName?.firstName;
            return (typeof firstName === 'string' && firstName.trim()) || 'there';
        } catch (err) {
            this.logger.warn(`Failed to fetch lead name for ${leadId}: ${err}`);
            return 'there';
        }
    }

    // Append an exchange to the lead's activity log. One activity record per
    // user/assistant turn. Safe to call in the background (we don't block the
    // stream on this).
    async logExchange(leadId: string, userText: string, aiText: string): Promise<void> {
        const summary = `User: ${userText}\n\nAI: ${aiText}`.slice(0, 4000);
        try {
            await this.platform.queryWithAuth<any>(
                `mutation($data: LeadActivityCreateInput!) {
                    createLeadActivity(data: $data) { id }
                }`,
                {
                    data: {
                        leadId,
                        activityType: 'NOTE_ADDED',
                        channel: 'SYSTEM',
                        summary,
                        occurredAt: new Date().toISOString(),
                        performedBy: 'Website Chat',
                        outcome: 'SUCCESSFUL',
                    },
                },
                this.auth,
            );
        } catch (err) {
            this.logger.warn(`Failed to log chat activity for lead ${leadId}: ${err}`);
        }
    }

    // Build a compact knowledge base of doctor/department info for the system
    // prompt. Cached for 5 min to avoid a doctors query on every chat.
    private async getKnowledgeBase(): Promise<string> {
        const now = Date.now();
        if (this.knowledgeBase && now - this.kbLoadedAt < this.kbTtlMs) {
            return this.knowledgeBase;
        }

        try {
            const doctors = await this.widget.getDoctors();
            const byDept = new Map<string, { name: string; visitingHours?: string; clinic?: string }[]>();
            for (const d of doctors) {
                const dept = (d.department ?? 'Other').replace(/_/g, ' ');
                if (!byDept.has(dept)) byDept.set(dept, []);
                byDept.get(dept)!.push({
                    name: d.name,
                    visitingHours: d.visitingHours,
                    clinic: d.clinic?.clinicName,
                });
            }

            const lines: string[] = ['DEPARTMENTS AND DOCTORS:'];
            for (const [dept, docs] of byDept) {
                lines.push(`\n${dept}:`);
                for (const doc of docs) {
                    const extras: string[] = [];
                    if (doc.visitingHours) extras.push(doc.visitingHours);
                    if (doc.clinic) extras.push(doc.clinic);
                    lines.push(`  - ${doc.name}${extras.length ? ` (${extras.join(' • ')})` : ''}`);
                }
            }
            this.knowledgeBase = lines.join('\n');
            this.kbLoadedAt = now;
        } catch (err) {
            this.logger.warn(`Failed to build widget KB: ${err}`);
            this.knowledgeBase = 'DEPARTMENTS AND DOCTORS: (unavailable)';
            this.kbLoadedAt = now;
        }
        return this.knowledgeBase;
    }

    async buildSystemPrompt(userName: string, selectedBranch: string | null): Promise<string> {
        const init = this.widget.getInitData();
        const kb = await this.getKnowledgeBase();

        // Branch context flips the tool-usage rules: no branch = must call
        // pick_branch first; branch set = always pass it to branch-aware
        // tools. We pre-render this block since the structure is dynamic
        // and the template just slots it in via {{branchContext}}.
        const branchContext = selectedBranch
            ? [
                  `CURRENT BRANCH: ${selectedBranch}`,
                  `The visitor is interested in the ${selectedBranch} branch. You MUST pass branch="${selectedBranch}"`,
                  'to list_departments, show_clinic_timings, show_doctors, and show_doctor_slots every time.',
              ].join('\n')
            : [
                  'BRANCH STATUS: NOT SET',
                  'The visitor has not picked a branch yet. Before calling list_departments, show_clinic_timings,',
                  'show_doctors, or show_doctor_slots, you MUST call pick_branch first so the visitor can choose.',
                  'Only skip this if the user asks a pure general question that does not need branch-specific data.',
              ].join('\n');

        return this.aiConfig.renderPrompt('widgetChat', {
            hospitalName: init.brand.name,
            userName,
            branchContext,
            knowledgeBase: kb,
        });
    }

    // Streams the assistant reply as an async iterable of UIMessageChunk-shaped
    // objects. The controller writes these as SSE `data: ${json}\n\n` lines
    // over the HTTP response. Tools return structured payloads the widget
    // frontend renders as generative-UI cards.
    async *streamReply(systemPrompt: string, messages: ModelMessage[]): AsyncGenerator<any> {
        const aiModel = this.buildAiModel();
        if (!aiModel) throw new Error('AI not configured');

        const platform = this.platform;
        const widgetSvc = this.widget;

        // Branch-matching now uses the doctor's full `clinics` array
        // (NormalizedDoctor) since one doctor can visit multiple
        // clinics under the post-rework data model. doctorMatchesBranch
        // returns true if ANY of their visit-slot clinics matches.
        const matchesBranch = (d: any, branch: string | undefined): boolean => {
            if (!branch) return true;
            const needle = branch.toLowerCase();
            const clinics: Array<{ clinicName: string }> = d.clinics ?? [];
            return clinics.some((c) => c.clinicName.toLowerCase().includes(needle));
        };

        const tools = {
            pick_branch: tool({
                description:
                    'Show the list of hospital branches so the visitor can pick which one they are interested in. Call this BEFORE any branch-sensitive tool (list_departments, show_clinic_timings, show_doctors, show_doctor_slots) when CURRENT BRANCH is NOT SET.',
                inputSchema: z.object({}),
                execute: async () => {
                    const doctors = await widgetSvc.getDoctors();
                    // Branches come from the union of all doctors'
                    // visit-slot clinics. Each (clinic × doctor) pair
                    // counts once toward that branch's doctor count;
                    // we use a Set on doctor ids to avoid double-
                    // counting the same doctor against the same branch
                    // when they have multiple slots there.
                    const byBranch = new Map<
                        string,
                        { doctorIds: Set<string>; departments: Set<string> }
                    >();
                    for (const d of doctors) {
                        for (const c of d.clinics ?? []) {
                            const name = c.clinicName?.trim();
                            if (!name) continue;
                            if (!byBranch.has(name)) {
                                byBranch.set(name, {
                                    doctorIds: new Set(),
                                    departments: new Set(),
                                });
                            }
                            const entry = byBranch.get(name)!;
                            if (d.id) entry.doctorIds.add(d.id);
                            if (d.department) entry.departments.add(String(d.department).replace(/_/g, ' '));
                        }
                    }
                    return {
                        branches: Array.from(byBranch.entries())
                            .sort(([a], [b]) => a.localeCompare(b))
                            .map(([name, { doctorIds, departments }]) => ({
                                name,
                                doctorCount: doctorIds.size,
                                departmentCount: departments.size,
                            })),
                    };
                },
            }),
            list_departments: tool({
                description:
                    'List the departments the hospital has. Use when the visitor asks what departments or specialities are available. Pass branch if CURRENT BRANCH is set.',
                inputSchema: z.object({
                    branch: z
                        .string()
                        .optional()
                        .describe('Branch name to filter by. Pass the CURRENT BRANCH when set.'),
                }),
                execute: async ({ branch }) => {
                    const doctors = await widgetSvc.getDoctors();
                    const filtered = doctors.filter((d: any) => matchesBranch(d, branch));
                    const deps = Array.from(
                        new Set(filtered.map((d: any) => d.department).filter(Boolean)),
                    ) as string[];
                    return {
                        branch: branch ?? null,
                        departments: deps.map(d => d.replace(/_/g, ' ')),
                    };
                },
            }),
            show_clinic_timings: tool({
                description:
                    'Show the clinic hours / visiting times for all departments with the doctors who visit during those hours. Use when the visitor asks about clinic timings, visiting hours, when the clinic is open, or what time a department is available. Pass branch if CURRENT BRANCH is set.',
                inputSchema: z.object({
                    branch: z
                        .string()
                        .optional()
                        .describe('Branch name to filter by. Pass the CURRENT BRANCH when set.'),
                }),
                execute: async ({ branch }) => {
                    const doctors = await widgetSvc.getDoctors();
                    const filtered = doctors.filter((d: any) => matchesBranch(d, branch));
                    const byDept = new Map<
                        string,
                        Array<{ name: string; hours: string; clinic: string | null }>
                    >();
                    for (const d of filtered) {
                        const dept = (d.department ?? 'Other').replace(/_/g, ' ');
                        if (!byDept.has(dept)) byDept.set(dept, []);
                        if (d.visitingHours) {
                            byDept.get(dept)!.push({
                                name: d.name,
                                hours: d.visitingHours,
                                clinic: d.clinic?.clinicName ?? null,
                            });
                        }
                    }
                    return {
                        branch: branch ?? null,
                        departments: Array.from(byDept.entries())
                            .filter(([, entries]) => entries.length > 0)
                            .map(([name, entries]) => ({ name, entries })),
                    };
                },
            }),
            show_doctors: tool({
                description:
                    'Show the list of doctors in a specific department with their visiting hours and clinic. Use when the visitor asks about doctors in a department. Pass branch if CURRENT BRANCH is set.',
                inputSchema: z.object({
                    department: z
                        .string()
                        .describe('Department name, e.g., "Cardiology", "ENT", "General Medicine".'),
                    branch: z
                        .string()
                        .optional()
                        .describe('Branch name to filter by. Pass the CURRENT BRANCH when set.'),
                }),
                execute: async ({ department, branch }) => {
                    const doctors = await widgetSvc.getDoctors();
                    const deptKey = department.toLowerCase().replace(/\s+/g, '').replace(/_/g, '');
                    const matches = doctors
                        .filter((d: any) => {
                            const key = String(d.department ?? '')
                                .toLowerCase()
                                .replace(/\s+/g, '')
                                .replace(/_/g, '');
                            return key.includes(deptKey) && matchesBranch(d, branch);
                        })
                        .map((d: any) => ({
                            id: d.id,
                            name: d.name,
                            specialty: d.specialty ?? null,
                            visitingHours: d.visitingHours ?? null,
                            clinic: d.clinic?.clinicName ?? null,
                        }));
                    return { department, branch: branch ?? null, doctors: matches };
                },
            }),
            show_doctor_slots: tool({
                description:
                    "Show today's available appointment slots for a specific doctor. Use when the visitor wants to see when a doctor is free or wants to book with a specific doctor. The date is always today — do NOT try to specify a date. Pass branch if CURRENT BRANCH is set to disambiguate doctors with the same name across branches.",
                inputSchema: z.object({
                    doctorName: z
                        .string()
                        .describe('Full name of the doctor, e.g., "Dr. Lakshmi Reddy".'),
                    branch: z
                        .string()
                        .optional()
                        .describe('Branch name to disambiguate. Pass the CURRENT BRANCH when set.'),
                }),
                execute: async ({ doctorName, branch }) => {
                    // Always use the server's current date. Never trust anything from
                    // the model here — older LLMs hallucinate their training-data
                    // "today" and return slots for the wrong day.
                    const targetDate = new Date().toISOString().slice(0, 10);
                    const doctors = await widgetSvc.getDoctors();
                    const scoped = doctors.filter((d: any) => matchesBranch(d, branch));
                    // Fuzzy match: lowercase + strip "Dr." prefix + collapse spaces.
                    const norm = (s: string) =>
                        s.toLowerCase().replace(/^dr\.?\s*/i, '').replace(/\s+/g, ' ').trim();
                    const target = norm(doctorName);
                    const doc =
                        scoped.find((d: any) => norm(d.name) === target) ??
                        scoped.find((d: any) => norm(d.name).includes(target)) ??
                        scoped.find((d: any) => target.includes(norm(d.name)));
                    if (!doc) {
                        return {
                            doctor: null,
                            date: targetDate,
                            slots: [],
                            error: `No doctor matching "${doctorName}"${branch ? ` at ${branch}` : ''} was found.`,
                        };
                    }
                    const slots = await widgetSvc.getSlots(doc.id, targetDate);
                    return {
                        doctor: {
                            id: doc.id,
                            name: doc.name,
                            department: doc.department ?? null,
                            clinic: doc.clinic?.clinicName ?? null,
                        },
                        date: targetDate,
                        slots,
                    };
                },
            }),
            suggest_booking: tool({
                description:
                    'Suggest that the visitor book an appointment. Use when the conversation is trending toward booking, the user has identified a concern, or asks "how do I book".',
                inputSchema: z.object({
                    reason: z.string().describe('Short reason why booking is a good next step.'),
                    department: z
                        .string()
                        .optional()
                        .describe('Suggested department, if known.'),
                }),
                execute: async ({ reason, department }) => {
                    return { reason, department: department ?? null };
                },
            }),
        };

        // Bookings / leads are not in scope for the AI — we only wire read/
        // suggest tools here. The CC agent/AP engineering team can book.
        void platform;

        const result = streamText({
            model: aiModel,
            system: systemPrompt,
            messages,
            tools,
            stopWhen: stepCountIs(4),
        });

        const uiStream = result.toUIMessageStream();
        for await (const chunk of uiStream) {
            yield chunk;
        }
    }
}