pediatric-ai-scribe-v3/client/src/pages/wellvisit/VisitNote.tsx

// ============================================================
// WELL VISIT — note generator (one of four sub-tabs).
// Picks up SSHADESS + Milestones carry-overs from sessionStorage
// (auto-set by Shadess.tsx and Milestones.tsx) so the user can
// flow byvisit → milestones → shadess → note without retyping.
// ============================================================

import { useEffect, useState } from 'react';
import { useMutation } from '@tanstack/react-query';
import { api, ApiError } from '@/lib/api';
import type { VisitNoteOk } from '@/shared/types';
import EncounterToolbar from '@/components/EncounterToolbar';
import OutputActions from '@/components/OutputActions';
import Recorder from '@/components/Recorder';

const TYPE = 'wellvisit' as const;
const input = 'w-full rounded-md border border-input bg-background px-3 py-2 text-sm';

export default function VisitNote() {
  const [label, setLabel] = useState('');
  const [patientAge, setPatientAge] = useState('');
  const [patientGender, setPatientGender] = useState('');
  const [visitAge, setVisitAge] = useState('');
  const [vitals, setVitals] = useState('');
  const [measurements, setMeasurements] = useState('');
  const [parentConcerns, setParentConcerns] = useState('');
  const [transcript, setTranscript] = useState('');
  const [interim, setInterim] = useState('');
  const [recError, setRecError] = useState<string | null>(null);
  const [shadess, setShadess] = useState('');
  const [milestones, setMilestones] = useState('');
  const [screenings, setScreenings] = useState('');
  const [vaccines, setVaccines] = useState('');
  const [byvisit, setByvisit] = useState('');
  const [noteStyle, setNoteStyle] = useState<'full' | 'short'>('full');
  const [result, setResult] = useState<string | null>(null);

  // On mount: pick up SSHADESS / Milestones / By-Visit carry-overs.
  useEffect(() => {
    try {
      const s = sessionStorage.getItem('wv-shadess-assessment');
      if (s) setShadess(s);
      const m = sessionStorage.getItem('wv-milestones-narrative');
      if (m) setMilestones(m);
      const b = sessionStorage.getItem('wv-byvisit-statuses');
      if (b) setByvisit(b);
      const va = sessionStorage.getItem('ped_visit_age');
      if (va) setVisitAge((cur) => cur || va);
    } catch { /* ignore */ }
  }, []);

  const generate = useMutation<VisitNoteOk, Error, Record<string, unknown>>({
    mutationFn: (body) => api.post<VisitNoteOk>('/api/well-visit/note', body),
    onSuccess: (data) => setResult(data.note),
  });

  function submit(e: React.FormEvent) {
    e.preventDefault();
    setResult(null);
    generate.mutate({
      patientAge, patientGender, visitAge,
      vitals, measurements, parentConcerns,
      transcript: (interim || transcript).trim(),
      shadessAssessment: shadess || undefined,
      // The byvisit summary is appended to screenings so the AI sees it.
      screenings: [screenings, byvisit].filter(Boolean).join('\n\n'),
      vaccines,
      // Milestones get folded into transcript context as a developmental block.
      physicianMemories: milestones ? '[DEVELOPMENTAL ASSESSMENT]\n' + milestones : undefined,
      noteStyle,
    });
  }

  const displayedTranscript = interim || transcript;

  return (
    <div className="space-y-4">
      <EncounterToolbar
        type={TYPE}
        label={label} setLabel={setLabel}
        transcript={transcript} generatedNote={result || ''}
        partialData={{ age: patientAge, gender: patientGender, visitAge, vitals, measurements, parentConcerns, shadess, milestones, byvisit, screenings, vaccines, noteStyle }}
        onLoad={(enc) => {
          setTranscript(enc.transcript || '');
          setInterim('');
          setResult(enc.generated_note || null);
          try {
            const pd = enc.partial_data ? JSON.parse(enc.partial_data) : null;
            if (pd?.age) setPatientAge(pd.age);
            if (pd?.gender) setPatientGender(pd.gender);
            if (pd?.visitAge) setVisitAge(pd.visitAge);
            if (pd?.vitals) setVitals(pd.vitals);
            if (pd?.measurements) setMeasurements(pd.measurements);
            if (pd?.parentConcerns) setParentConcerns(pd.parentConcerns);
            if (pd?.shadess) setShadess(pd.shadess);
            if (pd?.milestones) setMilestones(pd.milestones);
            if (pd?.byvisit) setByvisit(pd.byvisit);
            if (pd?.screenings) setScreenings(pd.screenings);
            if (pd?.vaccines) setVaccines(pd.vaccines);
            if (pd?.noteStyle) setNoteStyle(pd.noteStyle);
          } catch { /* ignore */ }
          setLabel(enc.label || '');
        }}
        onClear={() => {
          setTranscript(''); setInterim(''); setResult(null);
          setPatientAge(''); setPatientGender(''); setVisitAge('');
          setVitals(''); setMeasurements(''); setParentConcerns('');
          setShadess(''); setMilestones(''); setByvisit('');
          setScreenings(''); setVaccines(''); setNoteStyle('full');
        }}
      />

      <form onSubmit={submit} className="space-y-4">
        <div className="grid grid-cols-3 gap-3">
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Age</span>
            <input className={input} value={patientAge} onChange={(e) => setPatientAge(e.target.value)} />
          </label>
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Gender</span>
            <select className={input} value={patientGender} onChange={(e) => setPatientGender(e.target.value)}>
              <option value="">Select</option><option>Male</option><option>Female</option>
            </select>
          </label>
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Visit age</span>
            <input className={input} placeholder="e.g. 6 months" value={visitAge} onChange={(e) => setVisitAge(e.target.value)} />
          </label>
        </div>

        <div className="grid grid-cols-2 gap-3">
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Vital signs</span>
            <textarea className={input + ' min-h-[60px] font-mono text-xs'} value={vitals} onChange={(e) => setVitals(e.target.value)} />
          </label>
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Measurements / growth</span>
            <textarea className={input + ' min-h-[60px] font-mono text-xs'} value={measurements} onChange={(e) => setMeasurements(e.target.value)} />
          </label>
        </div>

        <label className="flex flex-col gap-1">
          <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Parent / patient concerns</span>
          <textarea className={input + ' min-h-[60px] text-sm'} value={parentConcerns} onChange={(e) => setParentConcerns(e.target.value)} />
        </label>

        <Recorder
          module="wellvisit"
          onTranscript={(text, meta) => {
            setTranscript((prev) => meta.appended ? (prev ? prev + ' ' + text : text) : text);
            setInterim('');
            setRecError(null);
          }}
          onInterim={(t) => setInterim(t ? (transcript ? transcript + ' ' + t : t) : '')}
          onError={(m) => setRecError(m)}
        />
        {recError && <div className="text-sm text-destructive">{recError}</div>}

        <label className="flex flex-col gap-1">
          <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Transcript / dictation</span>
          <textarea
            className={input + ' min-h-[160px] font-mono text-sm'}
            value={displayedTranscript}
            onChange={(e) => { setTranscript(e.target.value); setInterim(''); }}
            placeholder="Click Start recording, or type / paste."
          />
        </label>

        {(shadess || milestones || byvisit) && (
          <div className="rounded-lg border border-amber-200 bg-amber-50 dark:bg-amber-950/30 p-3 space-y-2">
            <div className="text-xs font-semibold text-amber-800 dark:text-amber-200">Carry-overs from other tabs (used in note generation)</div>
            {milestones && (
              <details className="text-xs">
                <summary className="cursor-pointer text-muted-foreground">Developmental assessment ({milestones.length} chars)</summary>
                <textarea className={input + ' mt-1 min-h-[80px] font-mono text-xs'} value={milestones} onChange={(e) => setMilestones(e.target.value)} />
              </details>
            )}
            {shadess && (
              <details className="text-xs">
                <summary className="cursor-pointer text-muted-foreground">SSHADESS assessment ({shadess.length} chars)</summary>
                <textarea className={input + ' mt-1 min-h-[80px] font-mono text-xs'} value={shadess} onChange={(e) => setShadess(e.target.value)} />
              </details>
            )}
            {byvisit && (
              <details className="text-xs">
                <summary className="cursor-pointer text-muted-foreground">By-visit-age statuses ({byvisit.split('\n').length} lines)</summary>
                <textarea className={input + ' mt-1 min-h-[80px] font-mono text-xs'} value={byvisit} onChange={(e) => setByvisit(e.target.value)} />
              </details>
            )}
          </div>
        )}

        <div className="grid grid-cols-2 gap-3">
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Screenings completed</span>
            <textarea className={input + ' min-h-[60px] text-xs'} value={screenings} onChange={(e) => setScreenings(e.target.value)} />
          </label>
          <label className="flex flex-col gap-1">
            <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Immunizations today</span>
            <textarea className={input + ' min-h-[60px] text-xs'} value={vaccines} onChange={(e) => setVaccines(e.target.value)} />
          </label>
        </div>

        <label className="flex flex-col gap-1">
          <span className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">Note style</span>
          <select className={input} value={noteStyle} onChange={(e) => setNoteStyle(e.target.value as 'full' | 'short')}>
            <option value="full">Full encounter note</option>
            <option value="short">Brief SOAP</option>
          </select>
        </label>

        {generate.error && <div className="text-sm text-destructive">{(generate.error as ApiError).message}</div>}

        <button
          type="submit"
          disabled={generate.isPending || (!patientAge.trim() && !visitAge.trim())}
          className="rounded-md bg-primary text-primary-foreground px-4 py-2 text-sm font-medium disabled:opacity-50"
        >
          {generate.isPending ? 'Generating…' : 'Generate Well Visit Note'}
        </button>
      </form>

      {result && (
        <section className="rounded-lg border border-border bg-card">
          <header className="px-4 py-2 border-b border-border bg-muted/40">
            <h2 className="text-sm font-semibold">Well Visit Note</h2>
          </header>
          <div className="p-4 whitespace-pre-wrap text-sm">{result}</div>
          <div className="px-4 pb-4">
            <OutputActions
              text={result}
              onUpdate={setResult}
              sourceContext={displayedTranscript}
              exportLabel="well-visit-note"
              exportType="well-visit"
            />
          </div>
        </section>
      )}
    </div>
  );
}