helix-engage-server/widget-src/src/chat.tsx

import { useState, useRef, useEffect } from 'preact/hooks';
import { startChatSession, streamChat } from './api';
import { readChatStream } from './chat-stream';
import { IconSpan } from './icon-span';
import { ChatToolWidget } from './chat-widgets';
import { useWidgetStore } from './store';
import type {
    BookingPrefill,
    ChatMessage,
    ChatPart,
    ChatTextPart,
    ChatToolPart,
} from './types';

const QUICK_ACTIONS = [
    'What departments do you have?',
    'Show me cardiologists',
    'Clinic timings',
    'How do I book?',
];

type ChatProps = {
    // Switches the widget to the Book tab. Chat-level handler that lives in
    // the parent so slot picks can seed bookingPrefill + swap tabs atomically.
    onRequestBooking: (prefill?: BookingPrefill) => void;
};

const textOf = (msg: ChatMessage): string =>
    msg.parts
        .filter((p): p is ChatTextPart => p.type === 'text')
        .map(p => p.text)
        .join('');

const updateMessage = (
    setMessages: (updater: (msgs: ChatMessage[]) => ChatMessage[]) => void,
    id: string,
    mutator: (msg: ChatMessage) => ChatMessage,
) => {
    setMessages(prev => prev.map(m => (m.id === id ? mutator(m) : m)));
};

const appendTextDelta = (parts: ChatPart[], delta: string, state: 'streaming' | 'done'): ChatPart[] => {
    const last = parts[parts.length - 1];
    if (last?.type === 'text') {
        return [...parts.slice(0, -1), { type: 'text', text: last.text + delta, state }];
    }
    return [...parts, { type: 'text', text: delta, state }];
};

const upsertToolPart = (
    parts: ChatPart[],
    toolCallId: string,
    update: Partial<ChatToolPart>,
    fallback: ChatToolPart,
): ChatPart[] => {
    const idx = parts.findIndex(p => p.type === 'tool' && p.toolCallId === toolCallId);
    if (idx >= 0) {
        const existing = parts[idx] as ChatToolPart;
        const merged: ChatToolPart = { ...existing, ...update };
        return [...parts.slice(0, idx), merged, ...parts.slice(idx + 1)];
    }
    return [...parts, { ...fallback, ...update }];
};

const genId = () => `m_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;

export const Chat = ({ onRequestBooking }: ChatProps) => {
    const {
        visitor,
        updateVisitor,
        leadId,
        setLeadId,
        setBookingPrefill,
        selectedBranch,
        setSelectedBranch,
    } = useWidgetStore();

    const [messages, setMessages] = useState<ChatMessage[]>([]);
    const [input, setInput] = useState('');
    const [loading, setLoading] = useState(false);
    const [formSubmitting, setFormSubmitting] = useState(false);
    const [formError, setFormError] = useState('');
    const scrollRef = useRef<HTMLDivElement>(null);

    useEffect(() => {
        if (scrollRef.current) {
            scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
        }
    }, [messages]);

    const submitLeadForm = async () => {
        const name = visitor.name.trim();
        const phone = visitor.phone.trim();
        if (!name || !phone) return;
        setFormSubmitting(true);
        setFormError('');
        try {
            const { leadId: newLeadId } = await startChatSession(name, phone);
            setLeadId(newLeadId);
        } catch {
            setFormError('Could not start chat. Please try again.');
        } finally {
            setFormSubmitting(false);
        }
    };

    const sendMessage = async (text: string, branchOverride?: string | null) => {
        if (!text.trim() || loading || !leadId) return;

        const userMsg: ChatMessage = {
            id: genId(),
            role: 'user',
            parts: [{ type: 'text', text: text.trim(), state: 'done' }],
        };
        const assistantId = genId();
        const assistantMsg: ChatMessage = { id: assistantId, role: 'assistant', parts: [] };

        const historyForBackend = [...messages, userMsg].map(m => ({
            role: m.role,
            content: textOf(m),
        }));

        setMessages(prev => [...prev, userMsg, assistantMsg]);
        setInput('');
        setLoading(true);

        // Branch can be provided explicitly to bypass the stale closure value
        // when the caller just set it (e.g., handleBranchPick immediately after
        // setSelectedBranch).
        const effectiveBranch =
            branchOverride !== undefined ? branchOverride : selectedBranch;

        try {
            const stream = await streamChat(leadId, historyForBackend, effectiveBranch);

            for await (const chunk of readChatStream(stream)) {
                switch (chunk.type) {
                    case 'text-delta':
                        if (typeof chunk.delta === 'string') {
                            updateMessage(setMessages, assistantId, m => ({
                                ...m,
                                parts: appendTextDelta(m.parts, chunk.delta, 'streaming'),
                            }));
                        }
                        break;
                    case 'text-end':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: m.parts.map(p =>
                                p.type === 'text' ? { ...p, state: 'done' } : p,
                            ),
                        }));
                        break;
                    case 'tool-input-start':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: upsertToolPart(
                                m.parts,
                                chunk.toolCallId,
                                { state: 'input-streaming', toolName: chunk.toolName },
                                {
                                    type: 'tool',
                                    toolCallId: chunk.toolCallId,
                                    toolName: chunk.toolName,
                                    state: 'input-streaming',
                                },
                            ),
                        }));
                        break;
                    case 'tool-input-available':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: upsertToolPart(
                                m.parts,
                                chunk.toolCallId,
                                { state: 'input-available', toolName: chunk.toolName, input: chunk.input },
                                {
                                    type: 'tool',
                                    toolCallId: chunk.toolCallId,
                                    toolName: chunk.toolName,
                                    state: 'input-available',
                                    input: chunk.input,
                                },
                            ),
                        }));
                        break;
                    case 'tool-output-available':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: upsertToolPart(
                                m.parts,
                                chunk.toolCallId,
                                { state: 'output-available', output: chunk.output },
                                {
                                    type: 'tool',
                                    toolCallId: chunk.toolCallId,
                                    toolName: 'unknown',
                                    state: 'output-available',
                                    output: chunk.output,
                                },
                            ),
                        }));
                        break;
                    case 'tool-output-error':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: upsertToolPart(
                                m.parts,
                                chunk.toolCallId,
                                { state: 'output-error', errorText: chunk.errorText },
                                {
                                    type: 'tool',
                                    toolCallId: chunk.toolCallId,
                                    toolName: 'unknown',
                                    state: 'output-error',
                                    errorText: chunk.errorText,
                                },
                            ),
                        }));
                        break;
                    case 'error':
                        updateMessage(setMessages, assistantId, m => ({
                            ...m,
                            parts: [
                                ...m.parts,
                                {
                                    type: 'text',
                                    text: 'Sorry, I encountered an error. Please try again.',
                                    state: 'done',
                                },
                            ],
                        }));
                        break;
                    default:
                        break;
                }
            }
        } catch {
            updateMessage(setMessages, assistantId, m => ({
                ...m,
                parts: [
                    ...m.parts,
                    {
                        type: 'text',
                        text: 'Sorry, I encountered an error. Please try again.',
                        state: 'done',
                    },
                ],
            }));
        } finally {
            setLoading(false);
        }
    };

    const handleSlotPick = (prefill: BookingPrefill) => {
        setBookingPrefill(prefill);
        onRequestBooking(prefill);
    };

    const handleBranchPick = (branch: string) => {
        // Store the selection so every subsequent request carries it, then
        // echo the visitor's choice as a user message so the AI re-runs the
        // branch-gated tool it was about to call. We pass branch explicitly
        // to sidestep the stale-closure selectedBranch inside sendMessage.
        setSelectedBranch(branch);
        sendMessage(`I'm interested in the ${branch} branch.`, branch);
    };

    // Pre-chat gate — only shown if we don't yet have an active lead. Name/phone
    // inputs bind to the shared store so anything typed here is immediately
    // available to the Book and Contact forms too.
    if (!leadId) {
        return (
            <div class="chat-intro">
                <div class="chat-empty-icon">
                    <IconSpan name="hand-wave" size={40} color="#f59e0b" />
                </div>
                <div class="chat-empty-title">Hi! How can we help?</div>
                <div class="chat-empty-text">
                    Share your name and phone so we can follow up if needed.
                </div>
                {formError && <div class="widget-error">{formError}</div>}
                <div class="widget-field">
                    <label class="widget-label">Full Name *</label>
                    <input
                        class="widget-input"
                        placeholder="Your name"
                        value={visitor.name}
                        onInput={(e: any) => updateVisitor({ name: e.target.value })}
                    />
                </div>
                <div class="widget-field">
                    <label class="widget-label">Phone Number *</label>
                    <input
                        class="widget-input"
                        placeholder="+91 9876543210"
                        value={visitor.phone}
                        onInput={(e: any) => updateVisitor({ phone: e.target.value })}
                        onKeyDown={(e: any) => e.key === 'Enter' && submitLeadForm()}
                    />
                </div>
                <button
                    class="widget-btn"
                    onClick={submitLeadForm}
                    disabled={!visitor.name.trim() || !visitor.phone.trim() || formSubmitting}
                >
                    {formSubmitting ? 'Starting…' : 'Start Chat'}
                </button>
            </div>
        );
    }

    return (
        <div style={{ display: 'flex', flexDirection: 'column', height: '100%' }}>
            <div class="chat-messages" ref={scrollRef}>
                {messages.length === 0 && (
                    <div class="chat-empty">
                        <div class="chat-empty-icon">
                            <IconSpan name="hand-wave" size={40} color="#f59e0b" />
                        </div>
                        <div class="chat-empty-title">
                            Hi {visitor.name.split(' ')[0] || 'there'}, how can we help?
                        </div>
                        <div class="chat-empty-text">
                            Ask about doctors, clinics, packages, or book an appointment.
                        </div>
                        <div class="quick-actions">
                            {QUICK_ACTIONS.map(q => (
                                <button key={q} class="quick-action" onClick={() => sendMessage(q)}>
                                    {q}
                                </button>
                            ))}
                        </div>
                    </div>
                )}
                {messages.map(msg => (
                    <MessageRow
                        key={msg.id}
                        msg={msg}
                        onDepartmentClick={dept => sendMessage(`Show me doctors in ${dept}`)}
                        onShowDoctorSlots={doctorName =>
                            sendMessage(`Show available appointments for ${doctorName}`)
                        }
                        onSuggestBooking={() => onRequestBooking()}
                        onPickSlot={handleSlotPick}
                        onPickBranch={handleBranchPick}
                    />
                ))}
            </div>
            <div class="chat-input-row">
                <input
                    class="widget-input chat-input"
                    placeholder="Type a message..."
                    value={input}
                    onInput={(e: any) => setInput(e.target.value)}
                    onKeyDown={(e: any) => e.key === 'Enter' && sendMessage(input)}
                    disabled={loading}
                />
                <button
                    class="chat-send"
                    onClick={() => sendMessage(input)}
                    disabled={loading || !input.trim()}
                    aria-label="Send message"
                >
                    <IconSpan name="paper-plane-top" size={16} color="#fff" />
                </button>
            </div>
        </div>
    );
};

type MessageRowProps = {
    msg: ChatMessage;
    onDepartmentClick: (dept: string) => void;
    onShowDoctorSlots: (doctorName: string) => void;
    onSuggestBooking: () => void;
    onPickSlot: (prefill: BookingPrefill) => void;
    onPickBranch: (branch: string) => void;
};

const MessageRow = ({
    msg,
    onDepartmentClick,
    onShowDoctorSlots,
    onSuggestBooking,
    onPickSlot,
    onPickBranch,
}: MessageRowProps) => {
    const isEmptyAssistant = msg.role === 'assistant' && msg.parts.length === 0;

    // If any tool parts exist, hide text parts from the same turn to avoid
    // models restating the widget's contents in prose.
    const hasToolParts = msg.parts.some(p => p.type === 'tool');
    const visibleParts = hasToolParts
        ? msg.parts.filter(p => p.type === 'tool')
        : msg.parts;

    return (
        <div class={`chat-msg ${msg.role}`}>
            <div class="chat-msg-stack">
                {isEmptyAssistant && (
                    <div class="chat-bubble">
                        <TypingDots />
                    </div>
                )}
                {visibleParts.map((part, i) => {
                    if (part.type === 'text') {
                        return (
                            <div key={i} class="chat-bubble">
                                {part.text || <TypingDots />}
                            </div>
                        );
                    }
                    return (
                        <ChatToolWidget
                            key={i}
                            part={part}
                            onDepartmentClick={onDepartmentClick}
                            onShowDoctorSlots={onShowDoctorSlots}
                            onSuggestBooking={onSuggestBooking}
                            onPickSlot={onPickSlot}
                            onPickBranch={onPickBranch}
                        />
                    );
                })}
            </div>
        </div>
    );
};

const TypingDots = () => (
    <span class="chat-typing-dots" aria-label="Assistant is typing">
        <span /><span /><span />
    </span>
);