maternal-app/maternal-web/app/(auth)/onboarding/page.tsx

'use client';

import { useState } from 'react';
import {
  Box,
  Stepper,
  Step,
  StepLabel,
  Button,
  Typography,
  Paper,
  TextField,
  Avatar,
  IconButton,
  Alert,
  CircularProgress,
  MenuItem,
  Card,
  CardActionArea,
  CardContent,
  Radio,
  RadioGroup,
  FormControlLabel,
  FormControl,
  Grid,
} from '@mui/material';
import { ArrowBack, ArrowForward, Check, Language, Straighten } from '@mui/icons-material';
import { motion, AnimatePresence } from 'framer-motion';
import { useRouter } from 'next/navigation';
import { useAuth } from '@/lib/auth/AuthContext';
import { childrenApi } from '@/lib/api/children';
import { useLocale, MeasurementSystem } from '@/hooks/useLocale';
import { useTranslation } from '@/hooks/useTranslation';
import { supportedLanguages } from '@/lib/i18n/config';
import { usersApi } from '@/lib/api/users';

const steps = ['Welcome', 'Language', 'Measurements', 'Add Child', 'Complete'];

export default function OnboardingPage() {
  const [activeStep, setActiveStep] = useState(0);
  const [selectedLanguage, setSelectedLanguage] = useState('en');
  const [selectedMeasurement, setSelectedMeasurement] = useState<MeasurementSystem>('metric');
  const [childName, setChildName] = useState('');
  const [childBirthDate, setChildBirthDate] = useState('');
  const [childGender, setChildGender] = useState<'male' | 'female' | 'other'>('other');
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState('');
  const router = useRouter();
  const { user, refreshUser } = useAuth();
  const { setLanguage, setMeasurementSystem } = useLocale();
  const { t } = useTranslation('onboarding');

  const handleNext = async () => {
    setError('');
    
    // Step 1: Save language preference
    if (activeStep === 1) {
      try {
        setLoading(true);
        await setLanguage(selectedLanguage);
        // Save to backend
        if (user?.id) {
          await usersApi.updatePreferences({
            language: selectedLanguage,
          });
        }
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } catch (err: any) {
        console.error('Failed to save language:', err);
        // Language was saved locally even if backend failed, so continue
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } finally {
        setLoading(false);
      }
      return;
    }

    // Step 2: Save measurement preference
    if (activeStep === 2) {
      try {
        setLoading(true);
        setMeasurementSystem(selectedMeasurement);
        // Save to backend
        if (user?.id) {
          await usersApi.updatePreferences({
            measurementUnit: selectedMeasurement,
          });
          await refreshUser();
        }
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } catch (err: any) {
        console.error('Failed to save measurement:', err);
        // Measurement was saved locally even if backend failed, so continue
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } finally {
        setLoading(false);
      }
      return;
    }

    // Step 3: Validate and save child data (Add Child)
    if (activeStep === 3) {
      if (!childName.trim() || !childBirthDate) {
        setError(t('child.name') + ' and ' + t('child.dateOfBirth') + ' are required');
        return;
      }

      const familyId = user?.families?.[0]?.familyId;
      if (!familyId) {
        console.error('No family found. User object:', JSON.stringify(user, null, 2));
        setError('Unable to create child profile. Your account setup is incomplete. Please contact support or try logging out and back in.');
        return;
      }

      try {
        setLoading(true);
        setError('');
        await childrenApi.createChild(familyId, {
          name: childName.trim(),
          birthDate: childBirthDate,
          gender: childGender,
        });
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } catch (err: any) {
        console.error('Failed to create child:', err);
        setError(err.response?.data?.message || 'Failed to save child. Please try again.');
      } finally {
        setLoading(false);
      }
      return;
    }

    if (activeStep === steps.length - 1) {
      // Complete onboarding
      router.push('/');
    } else {
      setActiveStep((prevActiveStep) => prevActiveStep + 1);
    }
  };

  const handleBack = () => {
    setActiveStep((prevActiveStep) => prevActiveStep - 1);
  };

  const handleSkip = () => {
    router.push('/');
  };

  return (
    <Box
      sx={{
        minHeight: '100vh',
        display: 'flex',
        flexDirection: 'column',
        px: 3,
        py: 4,
        background: 'linear-gradient(135deg, #FFE4E1 0%, #FFDAB9 100%)',
      }}
    >
      <Paper
        elevation={0}
        sx={{
          maxWidth: 600,
          mx: 'auto',
          width: '100%',
          p: 4,
          borderRadius: 4,
          background: 'rgba(255, 255, 255, 0.95)',
          backdropFilter: 'blur(10px)',
        }}
      >
        <Stepper activeStep={activeStep} sx={{ mb: 4 }}>
          {steps.map((label, index) => {
            let stepLabel = label;
            if (index === 0) stepLabel = t('welcome.title').split('!')[0];
            else if (index === 1) stepLabel = t('language.title');
            else if (index === 2) stepLabel = t('measurements.title');
            else if (index === 3) stepLabel = t('child.title');
            else if (index === 4) stepLabel = t('complete.title').split('!')[0];
            
            return (
              <Step key={label}>
                <StepLabel>{stepLabel}</StepLabel>
              </Step>
            );
          })}
        </Stepper>

        <AnimatePresence mode="wait">
          <motion.div
            key={activeStep}
            initial={{ opacity: 0, x: 20 }}
            animate={{ opacity: 1, x: 0 }}
            exit={{ opacity: 0, x: -20 }}
            transition={{ duration: 0.3 }}
          >
            {/* Step 0: Welcome */}
            {activeStep === 0 && (
              <Box sx={{ textAlign: 'center', py: 4 }}>
                <Typography variant="h4" gutterBottom fontWeight="600" color="primary.main">
                  {t('welcome.title')} 🎉
                </Typography>
                <Typography variant="body1" color="text.secondary" sx={{ mt: 2, mb: 4 }}>
                  {t('welcome.description')}
                </Typography>
                <Box sx={{ display: 'flex', gap: 2, justifyContent: 'center', flexWrap: 'wrap' }}>
                  <Paper sx={{ p: 2, flex: 1, minWidth: 150 }}>
                    <Typography variant="h6" fontWeight="600">📊</Typography>
                    <Typography variant="body2">{t('welcome.getStarted')}</Typography>
                  </Paper>
                  <Paper sx={{ p: 2, flex: 1, minWidth: 150 }}>
                    <Typography variant="h6" fontWeight="600">🤖</Typography>
                    <Typography variant="body2">AI Insights</Typography>
                  </Paper>
                  <Paper sx={{ p: 2, flex: 1, minWidth: 150 }}>
                    <Typography variant="h6" fontWeight="600">👨‍👩‍👧</Typography>
                    <Typography variant="body2">Family Sharing</Typography>
                  </Paper>
                </Box>
              </Box>
            )}

            {/* Step 1: Language Selection */}
            {activeStep === 1 && (
              <Box sx={{ py: 4 }}>
                <Box sx={{ display: 'flex', alignItems: 'center', mb: 2 }}>
                  <Language sx={{ fontSize: 40, color: 'primary.main', mr: 2 }} />
                  <Box>
                    <Typography variant="h5" fontWeight="600">
                      {t('language.title')}
                    </Typography>
                    <Typography variant="body2" color="text.secondary">
                      {t('language.subtitle')}
                    </Typography>
                  </Box>
                </Box>

                {error && (
                  <Alert severity="error" sx={{ mb: 2, borderRadius: 2 }}>
                    {error}
                  </Alert>
                )}

                <Grid container spacing={2} sx={{ mt: 2 }}>
                  {supportedLanguages.map((lang) => (
                    <Grid item xs={12} sm={6} key={lang.code}>
                      <Card
                        sx={{
                          border: selectedLanguage === lang.code ? '2px solid' : '1px solid',
                          borderColor: selectedLanguage === lang.code ? 'primary.main' : 'divider',
                          transition: 'all 0.2s',
                        }}
                      >
                        <CardActionArea
                          onClick={() => setSelectedLanguage(lang.code)}
                          sx={{ p: 2 }}
                        >
                          <Box sx={{ display: 'flex', alignItems: 'center', justifyContent: 'space-between' }}>
                            <Box>
                              <Typography variant="h6" fontWeight="600">
                                {lang.nativeName}
                              </Typography>
                              <Typography variant="body2" color="text.secondary">
                                {lang.name}
                              </Typography>
                            </Box>
                            <Radio
                              checked={selectedLanguage === lang.code}
                              value={lang.code}
                              name="language-radio"
                            />
                          </Box>
                        </CardActionArea>
                      </Card>
                    </Grid>
                  ))}
                </Grid>

                <Alert severity="info" sx={{ mt: 3, borderRadius: 2 }}>
                  {t('language.description')}
                </Alert>
              </Box>
            )}

            {/* Step 2: Measurement System */}
            {activeStep === 2 && (
              <Box sx={{ py: 4 }}>
                <Box sx={{ display: 'flex', alignItems: 'center', mb: 2 }}>
                  <Straighten sx={{ fontSize: 40, color: 'primary.main', mr: 2 }} />
                  <Box>
                    <Typography variant="h5" fontWeight="600">
                      {t('measurements.title')}
                    </Typography>
                    <Typography variant="body2" color="text.secondary">
                      {t('measurements.subtitle')}
                    </Typography>
                  </Box>
                </Box>

                {error && (
                  <Alert severity="error" sx={{ mb: 2, borderRadius: 2 }}>
                    {error}
                  </Alert>
                )}

                <Grid container spacing={3} sx={{ mt: 2 }}>
                  <Grid item xs={12} sm={6}>
                    <Card
                      sx={{
                        border: selectedMeasurement === 'metric' ? '2px solid' : '1px solid',
                        borderColor: selectedMeasurement === 'metric' ? 'primary.main' : 'divider',
                        height: '100%',
                      }}
                    >
                      <CardActionArea
                        onClick={() => setSelectedMeasurement('metric')}
                        sx={{ p: 3, height: '100%' }}
                      >
                        <Box sx={{ textAlign: 'center' }}>
                          <Radio
                            checked={selectedMeasurement === 'metric'}
                            value="metric"
                            name="measurement-radio"
                            sx={{ mb: 1 }}
                          />
                          <Typography variant="h6" fontWeight="600" gutterBottom>
                            {t('measurements.metric.title')}
                          </Typography>
                          <Typography variant="body2" color="text.secondary">
                            {t('measurements.metric.description')}
                          </Typography>
                        </Box>
                      </CardActionArea>
                    </Card>
                  </Grid>

                  <Grid item xs={12} sm={6}>
                    <Card
                      sx={{
                        border: selectedMeasurement === 'imperial' ? '2px solid' : '1px solid',
                        borderColor: selectedMeasurement === 'imperial' ? 'primary.main' : 'divider',
                        height: '100%',
                      }}
                    >
                      <CardActionArea
                        onClick={() => setSelectedMeasurement('imperial')}
                        sx={{ p: 3, height: '100%' }}
                      >
                        <Box sx={{ textAlign: 'center' }}>
                          <Radio
                            checked={selectedMeasurement === 'imperial'}
                            value="imperial"
                            name="measurement-radio"
                            sx={{ mb: 1 }}
                          />
                          <Typography variant="h6" fontWeight="600" gutterBottom>
                            {t('measurements.imperial.title')}
                          </Typography>
                          <Typography variant="body2" color="text.secondary">
                            {t('measurements.imperial.description')}
                          </Typography>
                        </Box>
                      </CardActionArea>
                    </Card>
                  </Grid>
                </Grid>

                <Alert severity="info" sx={{ mt: 3, borderRadius: 2 }}>
                  {t('measurements.description')}
                </Alert>
              </Box>
            )}

            {/* Step 3: Add Child */}
            {activeStep === 3 && (
              <Box sx={{ py: 4 }}>
                <Typography variant="h5" gutterBottom fontWeight="600">
                  {t('child.title')}
                </Typography>
                <Typography variant="body2" color="text.secondary" sx={{ mb: 3 }}>
                  {t('child.subtitle')}
                </Typography>

                {error && (
                  <Alert severity="error" sx={{ mb: 2, borderRadius: 2 }}>
                    {error}
                  </Alert>
                )}

                <TextField
                  fullWidth
                  label={t('child.name')}
                  value={childName}
                  onChange={(e) => setChildName(e.target.value)}
                  margin="normal"
                  required
                  disabled={loading}
                  InputProps={{
                    sx: { borderRadius: 3 },
                  }}
                />

                <TextField
                  fullWidth
                  label={t('child.dateOfBirth')}
                  type="date"
                  value={childBirthDate}
                  onChange={(e) => setChildBirthDate(e.target.value)}
                  margin="normal"
                  required
                  disabled={loading}
                  InputLabelProps={{
                    shrink: true,
                  }}
                  InputProps={{
                    sx: { borderRadius: 3 },
                  }}
                />

                <TextField
                  fullWidth
                  select
                  label={t('child.gender')}
                  value={childGender}
                  onChange={(e) => setChildGender(e.target.value as 'male' | 'female' | 'other')}
                  margin="normal"
                  disabled={loading}
                  InputProps={{
                    sx: { borderRadius: 3 },
                  }}
                >
                  <MenuItem value="male">{t('child.genders.male')}</MenuItem>
                  <MenuItem value="female">{t('child.genders.female')}</MenuItem>
                  <MenuItem value="other">{t('child.genders.preferNotToSay')}</MenuItem>
                </TextField>

                <Alert severity="info" sx={{ mt: 3, borderRadius: 2 }}>
                  {t('child.skipForNow')}
                </Alert>
              </Box>
            )}

            {/* Step 4: Complete */}
            {activeStep === 4 && (
              <Box sx={{ textAlign: 'center', py: 4 }}>
                <Avatar
                  sx={{
                    width: 80,
                    height: 80,
                    bgcolor: 'primary.main',
                    mx: 'auto',
                    mb: 3,
                  }}
                >
                  <Check sx={{ fontSize: 48 }} />
                </Avatar>
                <Typography variant="h5" gutterBottom fontWeight="600">
                  {t('complete.title')} 🎉
                </Typography>
                <Typography variant="body1" color="text.secondary" sx={{ mb: 4 }}>
                  {t('complete.description')}
                </Typography>

                <Paper sx={{ p: 3, bgcolor: 'primary.light', mb: 3 }}>
                  <Typography variant="body2" fontWeight="600" gutterBottom>
                    {t('complete.subtitle')}
                  </Typography>
                  <Typography variant="body2" align="left" component="div">
                    • Track your first feeding, sleep, or diaper change<br />
                    • Chat with our AI assistant for parenting tips<br />
                    • Explore insights and predictions based on your data
                  </Typography>
                </Paper>
              </Box>
            )}
          </motion.div>
        </AnimatePresence>

        <Box sx={{ display: 'flex', justifyContent: 'space-between', mt: 4 }}>
          <Button
            onClick={handleBack}
            disabled={activeStep === 0}
            startIcon={<ArrowBack />}
          >
            {t('navigation.back')}
          </Button>

          <Box sx={{ flex: 1 }} />

          {activeStep === 3 && (
            <Button onClick={handleSkip} sx={{ mr: 2 }}>
              {t('navigation.skip')}
            </Button>
          )}

          <Button
            variant="contained"
            onClick={handleNext}
            disabled={loading}
            endIcon={loading ? <CircularProgress size={20} /> : (activeStep === steps.length - 1 ? <Check /> : <ArrowForward />)}
          >
            {activeStep === steps.length - 1 ? t('complete.startTracking') : t('navigation.next')}
          </Button>
        </Box>
      </Paper>
    </Box>
  );
}