tia/src/app/(app)/onboarding/page.tsx

"use client";

import { useState, useEffect } from "react";
import { useRouter } from "next/navigation";
import { Button, Card, Input, Select } from "@/components/ui";

const IAP_SCHEDULE = [
  { name: "BCG",                  weeks: 0,   milestone: "At Birth" },
  { name: "OPV-0",                weeks: 0,   milestone: "At Birth" },
  { name: "HepB-1",               weeks: 0,   milestone: "At Birth" },
  { name: "OPV-1",                weeks: 6,   milestone: "6 Weeks" },
  { name: "Pentavalent-1",        weeks: 6,   milestone: "6 Weeks" },
  { name: "PCV-1",                weeks: 6,   milestone: "6 Weeks" },
  { name: "Rota-1",               weeks: 6,   milestone: "6 Weeks" },
  { name: "OPV-2",                weeks: 10,  milestone: "10 Weeks" },
  { name: "Pentavalent-2",        weeks: 10,  milestone: "10 Weeks" },
  { name: "PCV-2",                weeks: 10,  milestone: "10 Weeks" },
  { name: "Rota-2",               weeks: 10,  milestone: "10 Weeks" },
  { name: "OPV-3",                weeks: 14,  milestone: "14 Weeks" },
  { name: "Pentavalent-3",        weeks: 14,  milestone: "14 Weeks" },
  { name: "PCV-3",                weeks: 14,  milestone: "14 Weeks" },
  { name: "Rota-3",               weeks: 14,  milestone: "14 Weeks" },
  { name: "MR-1",                 weeks: 48,  milestone: "9 Months" },
  { name: "JE-1",                 weeks: 48,  milestone: "9 Months" },
  { name: "Vitamin A-1",          weeks: 48,  milestone: "9 Months" },
  { name: "OPV-4",                weeks: 48,  milestone: "9 Months" },
  { name: "MR-2",                 weeks: 96,  milestone: "18 Months" },
  { name: "JE-2",                 weeks: 96,  milestone: "18 Months" },
  { name: "DPT-Booster-1",        weeks: 96,  milestone: "18 Months" },
  { name: "Vitamin A-2",          weeks: 96,  milestone: "18 Months" },
  { name: "OPV-5",                weeks: 96,  milestone: "18 Months" },
  { name: "DPT-Booster-2",        weeks: 208, milestone: "4 Years" },
  { name: "Td",                   weeks: 208, milestone: "4 Years" },
];

function dueDate(birthDate: string, weeks: number) {
  const d = new Date(birthDate);
  d.setDate(d.getDate() + weeks * 7);
  return d.toISOString().split("T")[0];
}

function weeksSinceBirth(birthDate: string) {
  return Math.floor((Date.now() - new Date(birthDate).getTime()) / (1000 * 60 * 60 * 24 * 7));
}

type VaccState = { given: boolean; date: string; unknown: boolean };

export default function OnboardingPage() {
  const router = useRouter();
  const [step, setStep] = useState(1);
  const [loading, setLoading] = useState(false);
  const [saving, setSaving] = useState(false);
  const [checkingAuth, setCheckingAuth] = useState(true);
  const [newChildId, setNewChildId] = useState<string | null>(null);
  const [form, setForm] = useState({
    familyName: "",
    memberName: "",
    phone: "",
    childName: "",
    birthDate: "",
    sex: "" as "male" | "female" | "other",
  });
  const [vaccStates, setVaccStates] = useState<Record<string, VaccState>>({});

  useEffect(() => {
    async function checkAuth() {
      const res = await fetch("/api/auth/signin");
      const data = await res.json();
      if (!data.authenticated) {
        router.push("/login");
      } else if (data.familyId) {
        router.push("/home");
      }
      setCheckingAuth(false);
    }
    checkAuth();
  }, [router]);

  const dueVaccines = form.birthDate
    ? IAP_SCHEDULE.filter(v => v.weeks <= weeksSinceBirth(form.birthDate))
    : [];

  const milestones = [...new Set(dueVaccines.map(v => v.milestone))];

  const handleSubmit = async () => {
    setLoading(true);
    try {
      const res = await fetch("/api/onboarding", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(form),
      });
      const data = await res.json();
      if (data.success) {
        setNewChildId(data.childId);
        // Initialize vaccine states
        const initial: Record<string, VaccState> = {};
        dueVaccines.forEach(v => {
          initial[v.name] = { given: false, date: dueDate(form.birthDate, v.weeks), unknown: false };
        });
        setVaccStates(initial);
        setStep(3);
      }
    } catch (e) {
      console.error(e);
    }
    setLoading(false);
  };

  const handleVaccSave = async () => {
    setSaving(true);
    const entries = dueVaccines.map(v => {
      const s = vaccStates[v.name];
      return {
        vaccineName: v.name,
        scheduledDate: dueDate(form.birthDate, v.weeks),
        givenDate: s?.given ? s.date : undefined,
        status: s?.unknown ? "unknown" : s?.given ? "given" : "pending",
      };
    });
    try {
      await fetch("/api/vaccinations/bulk", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ childId: newChildId, entries }),
      });
    } catch (e) {
      console.error(e);
    }
    setSaving(false);
    router.push("/home");
  };

  if (checkingAuth) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-gradient-to-br from-rose-50 to-amber-50">
        <div className="text-gray-500">Loading...</div>
      </div>
    );
  }

  return (
    <div className="min-h-screen bg-gradient-to-br from-rose-50 to-amber-50 p-4">
      {/* Step indicators */}
      <div className="flex justify-center gap-2 pt-6 pb-4">
        {[1, 2, 3].map(s => (
          <div key={s} className={`h-1.5 w-10 rounded-full transition-colors ${step >= s ? "bg-rose-400" : "bg-gray-200"}`} />
        ))}
      </div>

      <div className="w-full max-w-md mx-auto">
        {step < 3 && (
          <div className="text-center mb-6">
            <span className="text-4xl">👶</span>
            <h1 className="text-2xl font-bold mt-3">Welcome to Tia</h1>
            <p className="text-gray-600">Let's set up your family</p>
          </div>
        )}

        <Card className="space-y-4" padding="lg">
          {/* Step 1 — Family info */}
          {step === 1 && (
            <>
              <h2 className="font-semibold text-gray-800">Your family</h2>
              <Input label="Family Name" type="text" value={form.familyName}
                onChange={(e) => setForm({ ...form, familyName: e.target.value })}
                placeholder="The Gupta Family" />
              <Input label="Your Name" type="text" value={form.memberName}
                onChange={(e) => setForm({ ...form, memberName: e.target.value })}
                placeholder="Mama" />
              <Input label="Phone Number (optional)" type="tel" value={form.phone}
                onChange={(e) => setForm({ ...form, phone: e.target.value })}
                placeholder="+91 98765 43210" />
              <p className="text-xs text-gray-400 -mt-2">For important reminders &amp; updates about your baby</p>
              <Button fullWidth size="lg" onClick={() => setStep(2)} disabled={!form.familyName || !form.memberName}>
                Next →
              </Button>
            </>
          )}

          {/* Step 2 — Child profile */}
          {step === 2 && (
            <>
              <h2 className="font-semibold text-gray-800">Your baby</h2>
              <Input label="Baby's Name" type="text" value={form.childName}
                onChange={(e) => setForm({ ...form, childName: e.target.value })}
                placeholder="Tia" />
              <Input label="Birth Date" type="date" value={form.birthDate}
                onChange={(e) => setForm({ ...form, birthDate: e.target.value })} />
              <div>
                <span className="text-sm font-medium text-gray-700">Sex</span>
                <div className="flex gap-2 mt-1">
                  {(["male", "female", "other"] as const).map((s) => (
                    <Button
                      key={s}
                      type="button"
                      fullWidth
                      variant={form.sex === s ? "primary" : "secondary"}
                      onClick={() => setForm({ ...form, sex: s })}
                      className="capitalize"
                    >
                      {s}
                    </Button>
                  ))}
                </div>
              </div>
              <div className="flex gap-2">
                <Button variant="secondary" fullWidth onClick={() => setStep(1)}>Back</Button>
                <Button fullWidth loading={loading}
                  disabled={!form.childName || !form.birthDate || !form.sex}
                  onClick={handleSubmit}>
                  Next →
                </Button>
              </div>
            </>
          )}

          {/* Step 3 — Vaccine history */}
          {step === 3 && (
            <>
              <div className="text-center mb-2">
                <span className="text-3xl">💉</span>
                <h2 className="font-bold text-lg mt-1">Vaccine History</h2>
                <p className="text-sm text-gray-500">
                  {form.childName} is {weeksSinceBirth(form.birthDate)} weeks old.
                  {dueVaccines.length > 0 ? ` Which vaccines have already been given?` : " No vaccines due yet."}
                </p>
              </div>

              {dueVaccines.length === 0 ? (
                <p className="text-center text-gray-400 text-sm py-4">All vaccines look up-to-date!</p>
              ) : (
                <div className="space-y-4 max-h-96 overflow-y-auto pr-1">
                  {milestones.map(milestone => (
                    <div key={milestone}>
                      <div className="text-xs font-semibold uppercase tracking-wider text-gray-400 mb-2">{milestone}</div>
                      <div className="space-y-2">
                        {dueVaccines.filter(v => v.milestone === milestone).map(v => {
                          const s = vaccStates[v.name] || { given: false, date: dueDate(form.birthDate, v.weeks), unknown: false };
                          return (
                            <div key={v.name} className="border rounded-xl p-3">
                              <div className="flex items-center justify-between gap-2">
                                <div className="flex items-center gap-2">
                                  <button
                                    type="button"
                                    onClick={() => setVaccStates(prev => ({ ...prev, [v.name]: { ...s, given: !s.given, unknown: false } }))}
                                    className={`w-5 h-5 rounded flex items-center justify-center text-xs transition-colors ${s.given ? "bg-green-500 text-white" : "border-2 border-gray-300"}`}
                                  >
                                    {s.given && "✓"}
                                  </button>
                                  <span className="text-sm font-medium">{v.name}</span>
                                </div>
                                <button
                                  type="button"
                                  onClick={() => setVaccStates(prev => ({ ...prev, [v.name]: { ...s, unknown: !s.unknown, given: false } }))}
                                  className={`text-xs px-2 py-0.5 rounded-full transition-colors ${s.unknown ? "bg-gray-200 text-gray-600" : "text-gray-400 hover:text-gray-600"}`}
                                >
                                  {s.unknown ? "Not sure" : "Not sure?"}
                                </button>
                              </div>
                              {s.given && (
                                <Input type="date" value={s.date} className="mt-2"
                                  onChange={e => setVaccStates(prev => ({ ...prev, [v.name]: { ...s, date: e.target.value } }))}
                                />
                              )}
                            </div>
                          );
                        })}
                      </div>
                    </div>
                  ))}
                </div>
              )}

              <div className="flex gap-2 pt-2">
                <Button variant="ghost" fullWidth onClick={() => router.push("/home")}>Skip for now</Button>
                <Button fullWidth loading={saving} onClick={handleVaccSave}>Save & Go Home</Button>
              </div>
            </>
          )}
        </Card>
      </div>
    </div>
  );
}