import { useState, useMemo, Fragment } from 'react';
import {
  Search,
  ChevronDown,
  ArrowUpDown,
  FileText,
  Stethoscope,
  Pill,
  AlertTriangle,
  Plus,
  X,
  Save,
  Building2,
  LayoutGrid,
  Edit2,
  Trash2,
} from 'lucide-react';
import { format, parseISO } from 'date-fns';
import { cn } from '../utils/helpers';
import { useHealthStore } from '../store/useHealthStore';
import { DatePickerInput } from '../components/DatePickerInput';
import { useActivePatientId } from '../hooks/useActivePatientId';
import type { APIHospitalVisit } from '../services/medical';

/* ── Types ─────────────────────────────────────────────────────────────────── */

interface Diagnosis {
  id: string;
  icdCode: string;
  diseaseName?: string;
  bodyRegion?: string;
  clinicalNotes?: string;
  practitioner?: string;
  clinic?: string;
  diagnosisDate?: string;
  severity: 'Mild' | 'Moderate' | 'Severe';
  isPrediagnosis: boolean;
}

interface Prescription {
  id: string;
  medicationName: string;
  dosage: string;
  frequency: string;
}

interface Visit {
  id: string;
  visitDate: string;
  hospital: string;
  department: string;
  doctor: string;
  dischargeSummary?: string;
  diagnoses: Diagnosis[];
  prescriptions: Prescription[];
}

function mapVisit(v: APIHospitalVisit): Visit {
  return {
    id: v._id,
    visitDate: v.visitDate,
    hospital: v.hospital,
    department: v.department || v.populatedDiagnoses?.[0]?.diagnosisClinic || 'General',
    doctor: v.doctorName || v.populatedDiagnoses?.[0]?.practitioner || 'Unknown',
    dischargeSummary: v.dischargeSummary || undefined,
    diagnoses: (v.populatedDiagnoses || []).map(d => ({
      id: d._id,
      icdCode: d.icdCode,
      diseaseName: d.diseaseName,
      bodyRegion: d.bodyRegion,
      clinicalNotes: d.notes || undefined,
      practitioner: d.practitioner,
      clinic: d.diagnosisClinic,
      severity: (() => {
        const s = (d.severity as string | undefined) || 'mild';
        return (s.charAt(0).toUpperCase() + s.slice(1)) as 'Mild' | 'Moderate' | 'Severe';
      })(),
      isPrediagnosis: d.isPrediagnosis || false,
    })),
    prescriptions: (v.populatedPrescriptions || []).map(p => ({
      id: p._id,
      medicationName: p.medicationName || 'Medication',
      dosage: p.dosage || '',
      frequency: p.frequency || '',
    })),
  };
}

/* ── Helpers ────────────────────────────────────────────────────────────────── */

type SortField = 'date' | 'hospital' | 'doctor';
type SortDirection = 'asc' | 'desc';

function initials(name: string) {
  const parts = name.trim().split(' ');
  if (parts.length >= 2) return parts.slice(-2).map(p => p[0]).join('').toUpperCase();
  return name.slice(0, 2).toUpperCase();
}

/* ── Component ──────────────────────────────────────────────────────────────── */

export function VisitsPage() {
  const [search, setSearch]               = useState('');
  const [showOnlyDiagnosed, setShowOnlyDiagnosed] = useState(false);
  const [expandedIds, setExpandedIds]     = useState<Set<string>>(new Set());
  const rawVisits = useHealthStore((s) => s._rawVisits);
  const loadedUserId = useHealthStore((s) => s._loadedUserId);
  const visits = useMemo(() => rawVisits.map(mapVisit), [rawVisits]);
  const isLoadingData = loadedUserId === null;
  const [sortField, setSortField]         = useState<SortField>('date');
  const [sortDir, setSortDir]             = useState<SortDirection>('desc');
  const activePatientId = useActivePatientId();
  const rawDiagnoses = useHealthStore((s) => s._rawDiagnoses);

  const addHospitalVisit = useHealthStore(s => s.addHospitalVisit);
  const editHospitalVisit = useHealthStore(s => s.editHospitalVisit);
  const removeHospitalVisit = useHealthStore(s => s.removeHospitalVisit);

  // Modal/Form States
  const [isModalOpen, setIsModalOpen] = useState(false);
  const [editingId, setEditingId] = useState<string | null>(null);
  const [processing, setProcessing] = useState(false);
  const [form, setForm] = useState({
    hospital: '',
    department: '',
    doctorName: '',
    visitDate: format(new Date(), 'yyyy-MM-dd'),
    admissionDate: '',
    dischargeDate: '',
    dischargeSummary: '',
    diagnoses: [] as string[]
  });

  function toggleExpand(id: string) {
    setExpandedIds(prev => {
      const next = new Set(prev);
      next.has(id) ? next.delete(id) : next.add(id);
      return next;
    });
  }

  function handleSort(field: SortField) {
    if (sortField === field) setSortDir(d => d === 'asc' ? 'desc' : 'asc');
    else { setSortField(field); setSortDir('desc'); }
  }

  const handleOpenAdd = () => {
    setEditingId(null);
    setForm({
      hospital: '',
      department: '',
      doctorName: '',
      visitDate: format(new Date(), 'yyyy-MM-dd'),
      admissionDate: '',
      dischargeDate: '',
      dischargeSummary: '',
      diagnoses: []
    });
    setIsModalOpen(true);
  };

  const handleOpenEdit = (v: APIHospitalVisit) => {
    setEditingId(v._id);
    setForm({
      hospital: v.hospital || '',
      department: v.department || '',
      doctorName: v.doctorName || '',
      visitDate: v.visitDate ? format(parseISO(v.visitDate), 'yyyy-MM-dd') : format(new Date(), 'yyyy-MM-dd'),
      admissionDate: v.admissionDate ? format(parseISO(v.admissionDate), 'yyyy-MM-dd') : '',
      dischargeDate: v.dischargeDate ? format(parseISO(v.dischargeDate), 'yyyy-MM-dd') : '',
      dischargeSummary: v.dischargeSummary || '',
      diagnoses: v.diagnoses || []
    });
    setIsModalOpen(true);
  };

  const handleDelete = async (id: string) => {
    if (!window.confirm("Are you sure you want to remove this visit record?")) return;
    try {
      await removeHospitalVisit(id);
    } catch (err) {
      alert("Failed to remove visit record");
    }
  };

  const handleSubmit = async () => {
    if (!activePatientId) return;
    if (!form.hospital || !form.visitDate) {
      alert("Please fill required fields (Hospital, Visit Date).");
      return;
    }

    setProcessing(true);
    try {
      const payload = {
        patientId: activePatientId,
        hospital: form.hospital,
        department: form.department || undefined,
        doctorName: form.doctorName || undefined,
        visitDate: new Date(form.visitDate).toISOString(),
        admissionDate: form.admissionDate ? new Date(form.admissionDate).toISOString() : undefined,
        dischargeDate: form.dischargeDate ? new Date(form.dischargeDate).toISOString() : undefined,
        dischargeSummary: form.dischargeSummary || undefined,
        diagnoses: form.diagnoses
      };

      if (editingId) {
        await editHospitalVisit(editingId, payload);
      } else {
        await addHospitalVisit(payload);
      }
      setIsModalOpen(false);
    } catch (err) {
      console.error("Operation failed", err);
      alert("Operation failed: " + (err as Error).message);
    } finally {
      setProcessing(false);
    }
  };

  const filtered = useMemo(() => {
    const q = search.toLowerCase();
    return visits
      .filter(v => {
        if (showOnlyDiagnosed && v.diagnoses.length === 0) return false;
        if (!q) return true;
        return [
          v.hospital, v.doctor, v.department,
          ...v.diagnoses.map(d => d.diseaseName || d.icdCode),
          ...v.prescriptions.map(p => p.medicationName),
        ].some(f => f.toLowerCase().includes(q));
      })
      .sort((a, b) => {
        let cmp = 0;
        if (sortField === 'date')     cmp = new Date(a.visitDate).getTime() - new Date(b.visitDate).getTime();
        if (sortField === 'hospital') cmp = a.hospital.localeCompare(b.hospital);
        if (sortField === 'doctor')   cmp = a.doctor.localeCompare(b.doctor);
        return sortDir === 'asc' ? cmp : -cmp;
      });
  }, [visits, search, showOnlyDiagnosed, sortField, sortDir]);

  const SortButton = ({ field, label }: { field: SortField; label: string }) => (
    <button
      onClick={() => handleSort(field)}
      className="flex items-center gap-1 group text-xs font-semibold text-gray-500 uppercase tracking-wider hover:text-gray-700"
    >
      <span>{label}</span>
      <ArrowUpDown className={cn(
        'w-3 h-3 transition-colors',
        sortField === field ? 'text-primary' : 'text-gray-300 group-hover:text-gray-400'
      )} />
    </button>
  );

  if (isLoadingData) {
    return (
      <div className="flex items-center justify-center py-32">
        <div className="animate-spin h-8 w-8 border-4 border-primary border-t-transparent rounded-full" />
      </div>
    );
  }

  return (
    <div className="pb-20">
      {/* Header + Controls */}
      <div className="flex flex-col md:flex-row justify-between items-start md:items-center mb-8 gap-4">
        <div>
          <h1 className="text-2xl font-semibold text-gray-900">My Hospital Visits</h1>
          <p className="text-gray-500 text-sm mt-1">Track your hospital appointments and related history</p>
        </div>
        <div className="flex items-center gap-2 w-full md:w-auto">
          <button
            onClick={() => setShowOnlyDiagnosed(v => !v)}
            className={cn(
              'px-4 py-2 rounded-lg text-sm font-medium border transition-colors whitespace-nowrap',
              showOnlyDiagnosed
                ? 'bg-primary text-white border-primary'
                : 'bg-white text-gray-700 border-gray-200 hover:bg-gray-50'
            )}
          >
            Show Diagnoses
          </button>
          <div className="relative flex-1 md:w-72">
            <Search className="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 w-4 h-4" />
            <input
              type="text"
              placeholder="Search hospital, doctor, diagnosis..."
              className="pl-10 pr-4 py-2 w-full bg-white border border-gray-200 rounded-lg text-sm focus:outline-none focus:ring-2 focus:ring-primary/20 transition-all"
              value={search}
              onChange={e => setSearch(e.target.value)}
            />
          </div>
          <button
            onClick={handleOpenAdd}
            className="flex items-center justify-center gap-2 px-4 py-2 rounded-lg bg-primary text-white text-sm font-bold shadow-lg shadow-primary/25 hover:-translate-y-0.5 active:translate-y-0 transition-all shrink-0"
          >
            <Plus className="w-4 h-4" />
            <span>Add Visit</span>
          </button>
        </div>
      </div>

      {/* Data Table */}
      <div className="bg-white rounded-xl border border-gray-200 shadow-sm overflow-hidden">
        <table className="w-full text-left table-fixed">
          <thead className="bg-gray-50 border-b border-gray-100">
            <tr>
              <th className="px-6 py-4 w-[140px]">
                <SortButton field="date" label="Date" />
              </th>
              <th className="px-6 py-4 w-[20%]">
                <SortButton field="hospital" label="Hospital" />
              </th>
              <th className="px-6 py-4 w-[16%]">
                <SortButton field="doctor" label="Doctor" />
              </th>
              <th className="px-6 py-4">
                <span className="text-xs font-semibold text-gray-500 uppercase tracking-wider">Diagnoses</span>
              </th>
              <th className="px-6 py-4 w-[16%]">
                <span className="text-xs font-semibold text-gray-500 uppercase tracking-wider">Prescriptions</span>
              </th>
              <th className="px-6 py-4 w-32" />
            </tr>
          </thead>
          <tbody className="divide-y divide-gray-100">
            {filtered.map(visit => {
              const isExpanded = expandedIds.has(visit.id);
              const visitDate = new Date(visit.visitDate);
              const rawVisit = rawVisits.find(rv => rv._id === visit.id);

              return (
                <Fragment key={visit.id}>
                  <tr
                    className={cn(
                      'hover:bg-gray-50/50 transition-all cursor-pointer group bg-white',
                      isExpanded && 'bg-primary/[0.02]'
                    )}
                  >
                    {/* Date */}
                    <td className="px-6 py-5 align-top" onClick={() => toggleExpand(visit.id)}>
                      <div className="text-sm font-bold text-gray-900">{format(visitDate, 'MMM d, yyyy')}</div>
                      <div className="text-[11px] font-medium text-gray-400 uppercase tracking-wider mt-1">{format(visitDate, 'EEEE')}</div>
                    </td>

                    {/* Hospital */}
                    <td className="px-6 py-5 align-top" onClick={() => toggleExpand(visit.id)}>
                      <p className="text-sm font-semibold text-gray-900 leading-snug">{visit.hospital}</p>
                      <p className="text-xs text-gray-500 mt-1 flex items-center gap-1">
                        <Building2 className="w-3 h-3" />
                        {visit.department}
                      </p>
                    </td>

                    {/* Doctor */}
                    <td className="px-6 py-5 align-top" onClick={() => toggleExpand(visit.id)}>
                      <div className="flex items-center gap-3">
                        <div className="w-9 h-9 rounded-full bg-primary/10 flex items-center justify-center text-[11px] font-bold text-primary border border-primary/20 shrink-0">
                          {initials(visit.doctor)}
                        </div>
                        <div className="min-w-0">
                          <span className="text-sm text-gray-700 font-semibold block truncate">Dr. {visit.doctor.replace(/^Dr\.?\s*/i, '')}</span>
                          <span className="text-[10px] text-gray-400 font-bold uppercase tracking-tighter">Attending Practitioner</span>
                        </div>
                      </div>
                    </td>

                    {/* Diagnoses */}
                    <td className="px-6 py-5 align-top" onClick={() => toggleExpand(visit.id)}>
                      {visit.diagnoses.length > 0 ? (
                        <div className="flex flex-wrap gap-1.5">
                          {visit.diagnoses.map(d => (
                            <div key={d.id} className="flex items-center gap-1.5 px-2 py-1 bg-gray-50 border border-gray-100 rounded-md">
                              <span className={cn(
                                'w-1.5 h-1.5 rounded-full shrink-0',
                                d.severity === 'Mild'     ? 'bg-emerald-400' :
                                d.severity === 'Moderate' ? 'bg-amber-400' :
                                                            'bg-red-500'
                              )} />
                              <span className="text-[11px] font-bold text-gray-600 leading-none">{d.diseaseName || d.icdCode}</span>
                              {d.isPrediagnosis && (
                                <AlertTriangle className="w-3 h-3 text-amber-500 shrink-0" />
                              )}
                            </div>
                          ))}
                        </div>
                      ) : (
                        <span className="text-sm text-gray-300">—</span>
                      )}
                    </td>

                    {/* Prescriptions */}
                    <td className="px-6 py-5 align-top" onClick={() => toggleExpand(visit.id)}>
                      {visit.prescriptions.length > 0 ? (
                        <div className="flex flex-wrap gap-1.5">
                          {visit.prescriptions.map(p => (
                            <div key={p.id} className="flex items-center gap-1.5 px-2 py-1 bg-violet-50/50 border border-violet-100 rounded-md">
                              <Pill className="w-3 h-3 text-violet-500 shrink-0" />
                              <span className="text-[11px] font-bold text-violet-700 leading-none">{p.medicationName}</span>
                            </div>
                          ))}
                        </div>
                      ) : (
                        <span className="text-sm text-gray-300">—</span>
                      )}
                    </td>

                    {/* Actions */}
                    <td className="px-6 py-5 text-right align-top">
                      <div className="flex items-center justify-end gap-1">
                        <div className="flex items-center gap-1">
                          <button
                            onClick={(e) => { e.stopPropagation(); if (rawVisit) handleOpenEdit(rawVisit); }}
                            className="p-1.5 text-gray-400 hover:text-primary hover:bg-primary/5 rounded-lg transition-all"
                            title="Edit Visit"
                          >
                            <Edit2 className="w-4 h-4" />
                          </button>
                          <button
                            onClick={(e) => { e.stopPropagation(); handleDelete(visit.id); }}
                            className="p-1.5 text-gray-400 hover:text-red-500 hover:bg-red-50 rounded-lg transition-all"
                            title="Remove Visit"
                          >
                            <Trash2 className="w-4 h-4" />
                          </button>
                        </div>
                        <button
                          onClick={(e) => { e.stopPropagation(); toggleExpand(visit.id); }}
                          className={cn(
                            'w-8 h-8 rounded-lg flex items-center justify-center transition-all duration-200',
                            isExpanded
                              ? 'bg-primary/10 text-primary'
                              : 'text-gray-400 hover:bg-gray-100'
                          )}
                        >
                          <ChevronDown className={cn('w-4 h-4 transition-transform duration-300', isExpanded && 'rotate-180')} />
                        </button>
                      </div>
                    </td>
                  </tr>

                  {/* Inline Expanded Details */}
                  {isExpanded && (
                    <tr className="bg-gray-50/50">
                      <td colSpan={6} className="px-6 py-8 border-t border-gray-100">
                        <div className="grid grid-cols-1 lg:grid-cols-3 gap-8 animate-in fade-in slide-in-from-top-2 duration-300">
                          {/* Column 1: Visit Details */}
                          <div className="bg-white rounded-2xl border border-gray-200 p-6 shadow-sm ring-1 ring-black/[0.02]">
                            <div className="flex items-center gap-2.5 mb-6">
                              <div className="w-8 h-8 rounded-lg bg-gray-100 flex items-center justify-center">
                                <FileText className="w-4 h-4 text-gray-500" />
                              </div>
                              <h4 className="text-xs font-bold text-gray-900 uppercase tracking-widest">Visit Details</h4>
                            </div>
                            <div className="space-y-5">
                              <div className="grid grid-cols-2 gap-4">
                                <div>
                                  <span className="text-[10px] font-bold text-gray-400 uppercase tracking-wider block mb-1">Hospital</span>
                                  <p className="text-sm text-gray-900 font-bold">{visit.hospital}</p>
                                </div>
                                <div>
                                  <span className="text-[10px] font-bold text-gray-400 uppercase tracking-wider block mb-1">Department</span>
                                  <p className="text-sm text-gray-900 font-semibold">{visit.department}</p>
                                </div>
                              </div>
                              <div>
                                <span className="text-[10px] font-bold text-gray-400 uppercase tracking-wider block mb-1">Attending Practitioner</span>
                                <p className="text-sm text-gray-900 font-bold">Dr. {visit.doctor.replace(/^Dr\.?\s*/i, '')}</p>
                              </div>
                              <div>
                                <span className="text-[10px] font-bold text-gray-400 uppercase tracking-wider block mb-1">Visit Date</span>
                                <p className="text-sm text-gray-900 font-semibold">
                                  {format(new Date(visit.visitDate), 'MMMM do, yyyy')}
                                </p>
                              </div>
                              {visit.dischargeSummary && (
                                <div className="pt-4 border-t border-gray-50">
                                  <span className="text-[10px] font-bold text-gray-400 uppercase tracking-wider block mb-2">Discharge Summary</span>
                                  <div className="bg-gray-50/50 rounded-xl p-4 border border-gray-100">
                                    <p className="text-xs text-gray-600 leading-relaxed italic">"{visit.dischargeSummary}"</p>
                                  </div>
                                </div>
                              )}
                            </div>
                          </div>

                          {/* Column 2: Diagnoses */}
                          <div className="bg-white rounded-2xl border border-gray-200 p-6 shadow-sm ring-1 ring-black/[0.02]">
                            <div className="flex items-center gap-2.5 mb-6">
                              <div className="w-8 h-8 rounded-lg bg-cyan-50 flex items-center justify-center">
                                <Stethoscope className="w-4 h-4 text-cyan-600" />
                              </div>
                              <h4 className="text-xs font-bold text-gray-900 uppercase tracking-widest">Diagnoses</h4>
                            </div>
                            {visit.diagnoses.length > 0 ? (
                              <div className="space-y-4">
                                {visit.diagnoses.map(d => (
                                  <div key={d.id} className="p-4 bg-gray-50/50 rounded-xl border border-gray-100 hover:border-cyan-200 transition-colors">
                                    <div className="flex items-start justify-between mb-3">
                                      <span className="inline-flex items-center px-2 py-0.5 rounded-md text-[10px] font-bold bg-cyan-100 text-cyan-700 border border-cyan-200">
                                        {d.icdCode}
                                      </span>
                                      <span className={cn(
                                        'px-2 py-0.5 rounded-full text-[9px] font-bold uppercase tracking-tighter',
                                        d.severity === 'Mild'     ? 'bg-emerald-100 text-emerald-700' :
                                        d.severity === 'Moderate' ? 'bg-amber-100 text-amber-700' :
                                                                    'bg-red-100 text-red-700'
                                      )}>
                                        {d.severity} Severity
                                      </span>
                                    </div>
                                    <p className="text-sm font-bold text-gray-900 mb-2 leading-snug">{d.diseaseName || d.icdCode}</p>
                                    
                                    {d.clinicalNotes && (
                                      <p className="text-[11px] text-gray-500 leading-relaxed line-clamp-2">{d.clinicalNotes}</p>
                                    )}
                                    
                                    <div className="mt-3 pt-3 border-t border-gray-200/50 flex items-center justify-between text-[10px] text-gray-400 font-bold uppercase tracking-wider">
                                      <span>{d.clinic}</span>
                                      {d.isPrediagnosis && (
                                        <span className="text-amber-600 flex items-center gap-1">
                                          <AlertTriangle className="w-3 h-3" /> Pre-diag
                                        </span>
                                      )}
                                    </div>
                                  </div>
                                ))}
                              </div>
                            ) : (
                              <div className="h-40 bg-gray-50/50 rounded-2xl border border-gray-100 border-dashed flex flex-col items-center justify-center">
                                <Stethoscope className="w-8 h-8 text-gray-200 mb-2" />
                                <p className="text-[11px] font-bold text-gray-400 uppercase">No diagnoses recorded.</p>
                              </div>
                            )}
                          </div>

                          {/* Column 3: Prescriptions */}
                          <div className="bg-white rounded-2xl border border-gray-200 p-6 shadow-sm ring-1 ring-black/[0.02]">
                            <div className="flex items-center gap-2.5 mb-6">
                              <div className="w-8 h-8 rounded-lg bg-violet-50 flex items-center justify-center">
                                <Pill className="w-4 h-4 text-violet-600" />
                              </div>
                              <h4 className="text-xs font-bold text-gray-900 uppercase tracking-widest">Linked Prescriptions</h4>
                            </div>
                            {visit.prescriptions.length > 0 ? (
                              <div className="space-y-3">
                                {visit.prescriptions.map(p => (
                                  <div key={p.id} className="group p-4 bg-violet-50/30 hover:bg-violet-50 rounded-xl border border-violet-100/50 transition-all">
                                    <div className="flex items-center gap-4">
                                      <div className="w-10 h-10 rounded-xl bg-white shadow-sm flex items-center justify-center text-violet-600 shrink-0 border border-violet-100">
                                        <Pill className="w-5 h-5" />
                                      </div>
                                      <div className="min-w-0">
                                        <p className="text-sm font-bold text-gray-900 truncate">{p.medicationName}</p>
                                        {(p.dosage || p.frequency) && (
                                          <p className="text-[11px] text-violet-600 font-bold mt-1">
                                            {[p.dosage, p.frequency].filter(Boolean).join(' • ')}
                                          </p>
                                        )}
                                      </div>
                                    </div>
                                  </div>
                                ))}
                              </div>
                            ) : (
                              <div className="h-40 bg-gray-50/50 rounded-2xl border border-gray-100 border-dashed flex flex-col items-center justify-center">
                                <Pill className="w-8 h-8 text-gray-200 mb-2" />
                                <p className="text-[11px] font-bold text-gray-400 uppercase">No prescriptions found.</p>
                              </div>
                            )}
                          </div>
                        </div>
                      </td>
                    </tr>
                  )}
                </Fragment>
              );
            })}
          </tbody>
        </table>

        {filtered.length === 0 && (
          <div className="flex flex-col items-center justify-center py-20 px-4">
            <div className="w-16 h-16 rounded-full bg-gray-100 flex items-center justify-center mb-4">
              <Search className="w-8 h-8 text-gray-400" />
            </div>
            <h3 className="text-lg font-semibold text-gray-900 mb-1">No visits found</h3>
            <p className="text-sm text-gray-500 text-center max-w-md">
              Try adjusting your filters or search terms.
            </p>
          </div>
        )}
      </div>

      {/* Modal */}
      {isModalOpen && (
        <div className="fixed inset-0 bg-black/60 z-[100] backdrop-blur-sm flex items-center justify-center p-4">
          <div className="bg-white rounded-2xl shadow-xl w-full max-w-lg overflow-hidden animate-in zoom-in-95 duration-200">
            <div className="p-4 border-b flex justify-between items-center bg-gray-50">
              <h3 className="font-bold text-gray-900">{editingId ? 'Edit Hospital Visit' : 'Add New Hospital Visit'}</h3>
              <button onClick={() => setIsModalOpen(false)} className="p-1 hover:bg-gray-200 rounded-md">
                <X className="w-5 h-5 text-gray-500" />
              </button>
            </div>
            <div className="p-6 space-y-4 max-h-[80vh] overflow-y-auto">
              <div className="grid grid-cols-2 gap-4">
                <div className="col-span-2">
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Hospital Name *</label>
                  <div className="relative">
                    <Building2 className="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 w-4 h-4" />
                    <input
                      type="text"
                      value={form.hospital}
                      onChange={e => setForm({ ...form, hospital: e.target.value })}
                      className="w-full pl-10 pr-3 py-2 rounded-lg border border-gray-200 focus:ring-2 focus:ring-primary/20"
                      placeholder="e.g. City General Hospital"
                    />
                  </div>
                </div>

                <div>
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Department</label>
                  <div className="relative">
                    <LayoutGrid className="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 w-4 h-4" />
                    <input
                      type="text"
                      value={form.department}
                      onChange={e => setForm({ ...form, department: e.target.value })}
                      className="w-full pl-10 pr-3 py-2 rounded-lg border border-gray-200 focus:ring-2 focus:ring-primary/20"
                      placeholder="e.g. Cardiology"
                    />
                  </div>
                </div>
                <div>
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Doctor Name</label>
                  <div className="relative">
                    <Stethoscope className="absolute left-3 top-1/2 -translate-y-1/2 text-gray-400 w-4 h-4" />
                    <input
                      type="text"
                      value={form.doctorName}
                      onChange={e => setForm({ ...form, doctorName: e.target.value })}
                      className="w-full pl-10 pr-3 py-2 rounded-lg border border-gray-200 focus:ring-2 focus:ring-primary/20"
                      placeholder="e.g. Dr. Smith"
                    />
                  </div>
                </div>

                <div className="col-span-2">
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Visit Date *</label>
                  <DatePickerInput value={form.visitDate} onChange={val => setForm({ ...form, visitDate: val })} className="w-full" />
                </div>

                <div>
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Admission Date</label>
                  <DatePickerInput value={form.admissionDate} onChange={val => setForm({ ...form, admissionDate: val })} className="w-full" />
                </div>
                <div>
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Discharge Date</label>
                  <DatePickerInput value={form.dischargeDate} onChange={val => setForm({ ...form, dischargeDate: val })} className="w-full" />
                </div>

                <div className="col-span-2">
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Related Diagnoses</label>
                  <div className="border border-gray-200 rounded-lg p-3 space-y-2 max-h-32 overflow-y-auto">
                    {rawDiagnoses.map(diag => (
                      <label key={diag._id} className="flex items-center gap-2 text-sm cursor-pointer hover:bg-gray-50 p-1 rounded">
                        <input
                          type="checkbox"
                          checked={form.diagnoses.includes(diag._id)}
                          onChange={e => {
                            if (e.target.checked) {
                              setForm({ ...form, diagnoses: [...form.diagnoses, diag._id] });
                            } else {
                              setForm({ ...form, diagnoses: form.diagnoses.filter(id => id !== diag._id) });
                            }
                          }}
                          className="rounded text-primary focus:ring-primary/20"
                        />
                        <span className="text-gray-700 truncate">{diag.disease?.name || diag.icdCode} ({format(parseISO(diag.diagnosisDate), 'MMM yyyy')})</span>
                      </label>
                    ))}
                  </div>
                </div>

                <div className="col-span-2">
                  <label className="block text-xs font-bold text-gray-500 uppercase tracking-wider mb-1.5">Discharge Summary</label>
                  <textarea
                    value={form.dischargeSummary}
                    onChange={e => setForm({ ...form, dischargeSummary: e.target.value })}
                    className="w-full px-3 py-2 rounded-lg border border-gray-200 focus:ring-2 focus:ring-primary/20 min-h-[100px]"
                    placeholder="Enter visit summary or notes..."
                  />
                </div>
              </div>

              <button
                onClick={handleSubmit}
                disabled={processing}
                className="w-full flex items-center justify-center gap-2 bg-primary text-white py-2.5 rounded-lg font-bold shadow-md hover:bg-primary/90 mt-2 transition-all disabled:opacity-50"
              >
                {processing ? <div className="w-5 h-5 border-2 border-white/30 border-t-white rounded-full animate-spin" /> : <Save className="w-4 h-4" />}
                {processing ? 'Processing...' : editingId ? 'Update Visit' : 'Add Visit'}
              </button>
            </div>
          </div>
        </div>
      )}
    </div>
  );
}
