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

'use client';

import { useState, useEffect } from 'react';
import {
  Box,
  Stepper,
  Step,
  StepLabel,
  Button,
  Typography,
  Paper,
  TextField,
  Avatar,
  IconButton,
  Alert,
  CircularProgress,
  MenuItem,
  Card,
  CardActionArea,
  CardContent,
  Radio,
  RadioGroup,
  FormControlLabel,
  FormControl,
  Grid,
  StepConnector,
  stepConnectorClasses,
  styled,
} 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';
import { useTheme } from '@mui/material/styles';
import { StepIconProps } from '@mui/material/StepIcon';
import { FamilySetupStep, ChildData } from '@/components/onboarding/FamilySetupStep';
import { familiesApi } from '@/lib/api/families';

const steps = ['Welcome', 'Preferences', 'Family Setup', 'Complete'];

// Custom connector for mobile-friendly stepper
const CustomConnector = styled(StepConnector)(({ theme }) => ({
  [`&.${stepConnectorClasses.active}`]: {
    [`& .${stepConnectorClasses.line}`]: {
      borderColor: theme.palette.primary.main,
    },
  },
  [`&.${stepConnectorClasses.completed}`]: {
    [`& .${stepConnectorClasses.line}`]: {
      borderColor: theme.palette.primary.main,
    },
  },
  [`& .${stepConnectorClasses.line}`]: {
    borderColor: theme.palette.divider,
    borderTopWidth: 2,
    borderRadius: 1,
  },
}));

// Custom step icon showing numbers
const CustomStepIconRoot = styled('div')<{ ownerState: { active?: boolean; completed?: boolean } }>(
  ({ theme, ownerState }) => ({
    display: 'flex',
    alignItems: 'center',
    justifyContent: 'center',
    width: 32,
    height: 32,
    borderRadius: '50%',
    backgroundColor: ownerState.completed
      ? theme.palette.primary.main
      : ownerState.active
      ? theme.palette.primary.main
      : theme.palette.grey[300],
    color: ownerState.active || ownerState.completed ? '#fff' : theme.palette.text.secondary,
    fontWeight: 600,
    fontSize: '0.875rem',
  })
);

function CustomStepIcon(props: StepIconProps) {
  const { active, completed, icon } = props;

  return (
    <CustomStepIconRoot ownerState={{ active, completed }}>
      {completed ? <Check sx={{ fontSize: 18 }} /> : icon}
    </CustomStepIconRoot>
  );
}

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 theme = useTheme();

  // Restore onboarding progress on mount
  useEffect(() => {
    const restoreProgress = async () => {
      try {
        if (user?.id) {
          const status = await usersApi.getOnboardingStatus();

          // If onboarding is completed, redirect to home
          if (status.completed) {
            router.push('/');
            return;
          }

          // If user has saved progress, restore it
          if (status.step > 0) {
            setActiveStep(status.step);

            if (status.data?.selectedLanguage) {
              setSelectedLanguage(status.data.selectedLanguage);
            }

            if (status.data?.selectedMeasurement) {
              setSelectedMeasurement(status.data.selectedMeasurement);
            }
          }
        }
      } catch (err) {
        console.error('Failed to restore onboarding progress:', err);
        // Continue with default state if restore fails
      }
    };

    restoreProgress();
  }, [user?.id, router]);

  const handleNext = async () => {
    setError('');

    // Step 1: Save preferences (language + measurement)
    if (activeStep === 1) {
      try {
        setLoading(true);
        await setLanguage(selectedLanguage);
        setMeasurementSystem(selectedMeasurement);

        // Save progress to backend
        if (user?.id) {
          await usersApi.updateOnboardingProgress({
            step: activeStep + 1,
            data: {
              selectedLanguage,
              selectedMeasurement,
            },
          });
          await refreshUser();
        }
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } catch (err: any) {
        console.error('Failed to save preferences:', err);
        // Preferences were saved locally even if backend failed, so continue
        setActiveStep((prevActiveStep) => prevActiveStep + 1);
      } finally {
        setLoading(false);
      }
      return;
    }

    // Step 2: Family Setup - This is now handled by FamilySetupStep component
    if (activeStep === 2) {
      // The component handles the logic internally via onCreateFamily or onJoinFamily
      return;
    }

    if (activeStep === steps.length - 1) {
      // Complete onboarding
      try {
        setLoading(true);
        if (user?.id) {
          await usersApi.completeOnboarding();
        }
        router.push('/');
      } catch (err) {
        console.error('Failed to mark onboarding complete:', err);
        // Navigate anyway since onboarding flow is done
        router.push('/');
      } finally {
        setLoading(false);
      }
    } else {
      setActiveStep((prevActiveStep) => prevActiveStep + 1);
    }
  };

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

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

  const handleCreateFamily = async (childData: ChildData) => {
    try {
      setLoading(true);
      setError('');

      // Create a new family for the user
      const createdFamily = await familiesApi.createFamily({
        name: `${user?.name || 'My'}'s Family`,
      });

      await refreshUser(); // Refresh to get new family membership

      // Add the first child to the new family
      await childrenApi.createChild(createdFamily.id, {
        name: childData.name.trim(),
        birthDate: childData.birthDate,
        gender: childData.gender,
      });

      setActiveStep((prev) => prev + 1);
    } catch (err: any) {
      console.error('Failed to create family:', err);
      setError(err.response?.data?.message || 'Failed to create family. Please try again.');
    } finally {
      setLoading(false);
    }
  };

  const handleJoinFamily = async (shareCode: string) => {
    try {
      setLoading(true);
      setError('');

      await familiesApi.joinFamily({ shareCode });
      await refreshUser(); // Refresh to get new family membership
      setActiveStep((prev) => prev + 1);
    } catch (err: any) {
      console.error('Failed to join family:', err);
      setError(err.response?.data?.message || 'Failed to join family. Please check the code and try again.');
    } finally {
      setLoading(false);
    }
  };

  return (
    <Box
      sx={{
        minHeight: '100vh',
        display: 'flex',
        flexDirection: 'column',
        px: 3,
        py: 4,
        background: `linear-gradient(135deg, ${theme.palette.primary.light} 0%, ${theme.palette.secondary.light} 100%)`,
      }}
    >
      <Paper
        elevation={0}
        sx={{
          maxWidth: 600,
          mx: 'auto',
          width: '100%',
          p: { xs: 3, sm: 4 },
          borderRadius: 4,
          background: 'rgba(255, 255, 255, 0.95)',
          backdropFilter: 'blur(10px)',
        }}
      >
        <Stepper
          activeStep={activeStep}
          alternativeLabel
          connector={<CustomConnector />}
          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 = 'Family Setup';
            else if (index === 4) stepLabel = t('complete.title').split('!')[0];

            // Only show label for active step on mobile
            const showLabel = activeStep === index;

            return (
              <Step key={label}>
                <StepLabel
                  StepIconComponent={CustomStepIcon}
                  sx={{
                    '& .MuiStepLabel-label': {
                      display: { xs: showLabel ? 'block' : 'none', sm: 'block' },
                      mt: 1,
                      fontSize: { xs: '0.75rem', sm: '0.875rem' },
                    },
                  }}
                >
                  {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: Preferences (Language + Measurements) */}
            {activeStep === 1 && (
              <Box sx={{ py: 4 }}>
                <Typography variant="h5" gutterBottom fontWeight="600">
                  {t('preferences.title') || 'Preferences'}
                </Typography>
                <Typography variant="body2" color="text.secondary" sx={{ mb: 3 }}>
                  {t('preferences.subtitle') || 'Set your language and measurement preferences'}
                </Typography>

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

                {/* Language Dropdown */}
                <TextField
                  select
                  fullWidth
                  label={t('language.title') || 'Language'}
                  value={selectedLanguage}
                  onChange={(e) => setSelectedLanguage(e.target.value)}
                  margin="normal"
                  sx={{ mb: 2 }}
                >
                  {supportedLanguages.map((lang) => (
                    <MenuItem key={lang.code} value={lang.code}>
                      {lang.nativeName} ({lang.name})
                    </MenuItem>
                  ))}
                </TextField>

                {/* Measurement Dropdown */}
                <TextField
                  select
                  fullWidth
                  label={t('measurements.title') || 'Measurement Units'}
                  value={selectedMeasurement}
                  onChange={(e) => setSelectedMeasurement(e.target.value as 'metric' | 'imperial')}
                  margin="normal"
                >
                  <MenuItem value="metric">
                    {t('measurements.metric.title') || 'Metric'} (kg, cm, °C)
                  </MenuItem>
                  <MenuItem value="imperial">
                    {t('measurements.imperial.title') || 'Imperial'} (lbs, inches, °F)
                  </MenuItem>
                </TextField>

                <Alert severity="info" sx={{ mt: 3, borderRadius: 2 }}>
                  {t('preferences.description') || 'You can change these settings later in your profile.'}
                </Alert>
              </Box>
            )}

            {/* Step 2: Family Setup */}
            {activeStep === 2 && (
              <FamilySetupStep
                onCreateFamily={handleCreateFamily}
                onJoinFamily={handleJoinFamily}
                loading={loading}
                error={error}
                t={t}
              />
            )}

            {/* Step 3: Complete */}
            {activeStep === 3 && (
              <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>
  );
}