import React, { useState, useMemo } from 'react';
import {
    Dialog, DialogTitle, DialogContent, DialogActions,
    Button, Stepper, Step, StepLabel, Box, TextField,
    Typography, IconButton, Select, MenuItem, FormControl,
    InputLabel, CircularProgress, Alert, Divider, Chip,
    Tooltip,
} from '@mui/material';
import { Add as AddIcon, Delete as DeleteIcon } from '@mui/icons-material';
import { useAuth } from '../../../../../../contexts/AuthContext';

interface TermInput {
    name: string;
    term_number: number;
    start_date: string;
    end_date: string;
    notes: string;
}

interface PeriodInput {
    code: string;
    name: string;
    start_time: string;
    end_time: string;
    period_type: 'lesson' | 'break' | 'registration' | 'offtimetable';
}

interface TermBreakInput {
    name: string;
    start_date: string;
    end_date: string;
}

interface WeekEntry {
    termName: string;
    termNumber: number;
    weekNumber: number;
    startDate: string;
    cycle: 'A' | 'B';
}

export interface SchoolInfo {
    name: string;
    urn: string;
    website: string;
    address: Record<string, string>;
    headteacher: string;
    term_dates_url: string;
    staff_list_url: string;
}

const DEFAULT_TERMS: TermInput[] = [
    { name: 'Autumn', term_number: 1, start_date: '2025-09-03', end_date: '2025-12-19', notes: '' },
    { name: 'Spring', term_number: 2, start_date: '2026-01-06', end_date: '2026-04-01', notes: '' },
    { name: 'Summer', term_number: 3, start_date: '2026-04-22', end_date: '2026-07-22', notes: '' },
];

const DEFAULT_TERM_BREAKS: TermBreakInput[] = [
    { name: 'Christmas Break', start_date: '2025-12-22', end_date: '2026-01-02' },
    { name: 'Easter Break', start_date: '2026-04-06', end_date: '2026-04-17' },
];

const DEFAULT_PERIODS: PeriodInput[] = [
    { code: 'REG', name: 'Registration', start_time: '08:45', end_time: '09:00', period_type: 'registration' },
    { code: 'P1', name: 'Period 1', start_time: '09:00', end_time: '10:00', period_type: 'lesson' },
    { code: 'P2', name: 'Period 2', start_time: '10:00', end_time: '11:00', period_type: 'lesson' },
    { code: 'BREAK', name: 'Break', start_time: '11:00', end_time: '11:20', period_type: 'break' },
    { code: 'P3', name: 'Period 3', start_time: '11:20', end_time: '12:20', period_type: 'lesson' },
    { code: 'P4', name: 'Period 4', start_time: '12:20', end_time: '13:20', period_type: 'lesson' },
    { code: 'LUNCH', name: 'Lunch', start_time: '13:20', end_time: '14:05', period_type: 'break' },
    { code: 'P5', name: 'Period 5', start_time: '14:05', end_time: '15:05', period_type: 'lesson' },
];

// Mirror of Python _academic_weeks: yields {n, monday} for each Mon block overlapping term.
function getAcademicWeeksForTerm(
    termStartStr: string,
    termEndStr: string,
): { n: number; monday: string }[] {
    if (!termStartStr || !termEndStr) return [];
    const termStart = new Date(termStartStr + 'T00:00:00');
    const termEnd = new Date(termEndStr + 'T00:00:00');
    const dayOfWeek = termStart.getDay(); // 0=Sun
    const daysFromMonday = dayOfWeek === 0 ? 6 : dayOfWeek - 1;
    const current = new Date(termStart);
    current.setDate(current.getDate() - daysFromMonday);
    if (current < termStart) current.setDate(current.getDate() + 7);
    const weeks: { n: number; monday: string }[] = [];
    let n = 1;
    while (current <= termEnd) {
        weeks.push({ n, monday: current.toISOString().slice(0, 10) });
        current.setDate(current.getDate() + 7);
        n++;
    }
    return weeks;
}

function computeWeekEntries(terms: TermInput[]): WeekEntry[] {
    const entries: WeekEntry[] = [];
    for (const term of terms) {
        if (!term.start_date || !term.end_date) continue;
        const termWeeks = getAcademicWeeksForTerm(term.start_date, term.end_date);
        for (const { n, monday } of termWeeks) {
            // default A/B: odd week within term = A, even = B (mirrors backend)
            const cycle: 'A' | 'B' = n % 2 === 1 ? 'A' : 'B';
            entries.push({
                termName: term.name,
                termNumber: term.term_number,
                weekNumber: n,
                startDate: monday,
                cycle,
            });
        }
    }
    return entries;
}

interface Props {
    open: boolean;
    onClose: () => void;
    onComplete: () => void;
    apiBase: string;
    schoolInfo: SchoolInfo;
}

export function SchoolCalendarWizard({ open, onClose, onComplete, apiBase, schoolInfo }: Props) {
    const { accessToken } = useAuth();
    const [step, setStep] = useState(0);
    const [saving, setSaving] = useState(false);
    const [error, setError] = useState<string | null>(null);

    // Step 0
    const [headteacher, setHeadteacher] = useState(schoolInfo.headteacher || '');
    const [termDatesUrl, setTermDatesUrl] = useState(schoolInfo.term_dates_url || '');
    const [staffListUrl, setStaffListUrl] = useState(schoolInfo.staff_list_url || '');

    // Step 1
    const [yearStart, setYearStart] = useState('2025-09-01');
    const [yearEnd, setYearEnd] = useState('2026-07-31');
    const [terms, setTerms] = useState<TermInput[]>(DEFAULT_TERMS);

    // Step 2
    const [termBreaks, setTermBreaks] = useState<TermBreakInput[]>(DEFAULT_TERM_BREAKS);

    // Step 3 (periods)
    const [periods, setPeriods] = useState<PeriodInput[]>(DEFAULT_PERIODS);

    // Step 4 (week cycles) — computed from terms, user can override
    const [weekEntries, setWeekEntries] = useState<WeekEntry[]>([]);

    // ── Term helpers ─────────────────────────────────────────────────────────

    const addTerm = () => setTerms(prev => [...prev, {
        name: `Term ${prev.length + 1}`,
        term_number: prev.length + 1,
        start_date: '',
        end_date: '',
        notes: '',
    }]);

    const removeTerm = (i: number) => setTerms(prev =>
        prev.filter((_, idx) => idx !== i).map((t, idx) => ({ ...t, term_number: idx + 1 }))
    );

    const updateTerm = (i: number, field: keyof TermInput, value: string) =>
        setTerms(prev => prev.map((t, idx) => idx === i ? { ...t, [field]: value } : t));

    // ── Term break helpers ───────────────────────────────────────────────────

    const addTermBreak = () => setTermBreaks(prev => [...prev, {
        name: `Break ${prev.length + 1}`,
        start_date: '',
        end_date: '',
    }]);

    const removeTermBreak = (i: number) => setTermBreaks(prev => prev.filter((_, idx) => idx !== i));

    const updateTermBreak = (i: number, field: keyof TermBreakInput, value: string) =>
        setTermBreaks(prev => prev.map((b, idx) => idx === i ? { ...b, [field]: value } : b));

    // ── Period helpers ───────────────────────────────────────────────────────

    const addPeriod = () => setPeriods(prev => [...prev, {
        code: `P${prev.length + 1}`,
        name: `Period ${prev.length + 1}`,
        start_time: '',
        end_time: '',
        period_type: 'lesson',
    }]);

    const removePeriod = (i: number) => setPeriods(prev => prev.filter((_, idx) => idx !== i));

    const updatePeriod = (i: number, field: keyof PeriodInput, value: string) =>
        setPeriods(prev => prev.map((p, idx) => idx === i ? { ...p, [field]: value } : p));

    // ── Week cycle helpers ───────────────────────────────────────────────────

    const toggleWeekCycle = (idx: number) =>
        setWeekEntries(prev => prev.map((w, i) =>
            i === idx ? { ...w, cycle: w.cycle === 'A' ? 'B' : 'A' } : w
        ));

    const resetWeekCycles = () =>
        setWeekEntries(computeWeekEntries(terms));

    // ── Navigation ───────────────────────────────────────────────────────────

    const goToStep = (next: number) => {
        // When entering step 4, compute week entries from current terms
        if (next === 4) {
            setWeekEntries(computeWeekEntries(terms));
        }
        setStep(next);
    };

    // ── Step 0: save school details ──────────────────────────────────────────

    const handleSaveSchoolInfo = async () => {
        if (!accessToken) return;
        setSaving(true);
        setError(null);
        try {
            const res = await fetch(`${apiBase}/school/info`, {
                method: 'PATCH',
                headers: { Authorization: `Bearer ${accessToken}`, 'Content-Type': 'application/json' },
                body: JSON.stringify({ headteacher, term_dates_url: termDatesUrl, staff_list_url: staffListUrl }),
            });
            const data = await res.json();
            if (data.status === 'ok') {
                setStep(1);
            } else {
                setError(data.message || 'Failed to save school info');
            }
        } catch (e: any) {
            setError(e.message);
        } finally {
            setSaving(false);
        }
    };

    // ── Step 4: final save ───────────────────────────────────────────────────

    const handleSaveCalendar = async () => {
        if (!accessToken) return;
        setSaving(true);
        setError(null);
        try {
            const payload = {
                year_start: yearStart,
                year_end: yearEnd,
                terms: terms.map(t => ({
                    name: t.name,
                    term_number: t.term_number,
                    start_date: t.start_date,
                    end_date: t.end_date,
                    notes: t.notes || undefined,
                })),
                periods,
                term_breaks: termBreaks.filter(b => b.name && b.start_date && b.end_date),
                week_cycles: weekEntries.map(w => ({
                    term_number: w.termNumber,
                    week_number: w.weekNumber,
                    cycle: w.cycle,
                })),
            };

            const setupRes = await fetch(`${apiBase}/timetable/setup`, {
                method: 'POST',
                headers: { Authorization: `Bearer ${accessToken}`, 'Content-Type': 'application/json' },
                body: JSON.stringify(payload),
            });
            const setupData = await setupRes.json();
            if (setupData.status !== 'ok') {
                setError(setupData.message || 'Calendar setup failed');
                return;
            }

            // Materialize academic_periods rows
            const matRes = await fetch(`${apiBase}/timetable/materialize-periods`, {
                method: 'POST',
                headers: { Authorization: `Bearer ${accessToken}`, 'Content-Type': 'application/json' },
            });
            const matData = await matRes.json();
            if (matData.status !== 'ok') {
                // Non-fatal: log and continue
                console.warn('Materialize periods returned:', matData);
            }

            onComplete();
            handleClose();
        } catch (e: any) {
            setError(e.message);
        } finally {
            setSaving(false);
        }
    };

    const handleClose = () => {
        setStep(0);
        setError(null);
        onClose();
    };

    const addr = schoolInfo.address || {};
    const addressStr = [addr.street, addr.town, addr.county, addr.postcode].filter(Boolean).join(', ');

    const STEPS = ['School Details', 'Terms', 'Term Breaks', 'Daily Periods', 'Week Cycles'];

    // Group week entries by term for display
    const weeksByTerm = useMemo(() => {
        const map: Record<string, WeekEntry[]> = {};
        for (const w of weekEntries) {
            const key = w.termName;
            if (!map[key]) map[key] = [];
            map[key].push(w);
        }
        return map;
    }, [weekEntries]);

    return (
        <Dialog open={open} onClose={handleClose} maxWidth="md" fullWidth>
            <DialogTitle sx={{ pb: 1 }}>Set Up School Calendar</DialogTitle>
            <Box sx={{ px: 3 }}>
                <Stepper activeStep={step} sx={{ mb: 2 }}>
                    {STEPS.map(label => <Step key={label}><StepLabel>{label}</StepLabel></Step>)}
                </Stepper>
            </Box>

            <DialogContent>
                {error && <Alert severity="error" sx={{ mb: 2 }}>{error}</Alert>}

                {/* ── Step 0: School Details ───────────────────────────────── */}
                {step === 0 && (
                    <Box>
                        <Typography variant="subtitle2" sx={{ mb: 1.5 }}>School Information</Typography>
                        <Box sx={{ mb: 2, p: 1.5, bgcolor: 'action.hover', borderRadius: 1 }}>
                            <Typography variant="body2" sx={{ fontWeight: 600 }}>{schoolInfo.name || '—'}</Typography>
                            {schoolInfo.urn && (
                                <Typography variant="caption" sx={{ color: 'text.secondary' }}>URN: {schoolInfo.urn}</Typography>
                            )}
                            {addressStr && (
                                <Typography variant="caption" sx={{ display: 'block', color: 'text.secondary' }}>{addressStr}</Typography>
                            )}
                            {schoolInfo.website && (
                                <Typography variant="caption" sx={{ display: 'block', color: 'text.secondary' }}>{schoolInfo.website}</Typography>
                            )}
                        </Box>
                        <Divider sx={{ mb: 2 }} />
                        <Typography variant="subtitle2" sx={{ mb: 1.5 }}>Additional Details</Typography>
                        <Box sx={{ display: 'flex', flexDirection: 'column', gap: 2 }}>
                            <TextField
                                label="Headteacher"
                                value={headteacher}
                                onChange={e => setHeadteacher(e.target.value)}
                                size="small"
                                fullWidth
                                placeholder="e.g. Mr J Smith"
                            />
                            <TextField
                                label="Term Dates URL"
                                value={termDatesUrl}
                                onChange={e => setTermDatesUrl(e.target.value)}
                                size="small"
                                fullWidth
                                placeholder="Link to term dates page on school website"
                            />
                            <TextField
                                label="Staff List URL"
                                value={staffListUrl}
                                onChange={e => setStaffListUrl(e.target.value)}
                                size="small"
                                fullWidth
                                placeholder="Link to staff list page on school website"
                            />
                        </Box>
                    </Box>
                )}

                {/* ── Step 1: Terms ────────────────────────────────────────── */}
                {step === 1 && (
                    <Box>
                        <Typography variant="subtitle2" sx={{ mb: 1.5 }}>School Year</Typography>
                        <Box sx={{ display: 'flex', gap: 2, mb: 3 }}>
                            <TextField label="Year Start" type="date" value={yearStart}
                                onChange={e => setYearStart(e.target.value)}
                                size="small" InputLabelProps={{ shrink: true }} />
                            <TextField label="Year End" type="date" value={yearEnd}
                                onChange={e => setYearEnd(e.target.value)}
                                size="small" InputLabelProps={{ shrink: true }} />
                        </Box>
                        <Divider sx={{ mb: 2 }} />
                        <Box sx={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', mb: 1.5 }}>
                            <Typography variant="subtitle2">Terms</Typography>
                            <Button size="small" startIcon={<AddIcon />} onClick={addTerm}>Add Term</Button>
                        </Box>
                        {terms.map((term, i) => (
                            <Box key={i} sx={{ mb: 2 }}>
                                <Box sx={{ display: 'flex', gap: 1.5, mb: 0.75, alignItems: 'center' }}>
                                    <TextField label="Term Name" value={term.name}
                                        onChange={e => updateTerm(i, 'name', e.target.value)}
                                        size="small" sx={{ width: 140 }} />
                                    <TextField label="Start Date" type="date" value={term.start_date}
                                        onChange={e => updateTerm(i, 'start_date', e.target.value)}
                                        size="small" InputLabelProps={{ shrink: true }} />
                                    <TextField label="End Date" type="date" value={term.end_date}
                                        onChange={e => updateTerm(i, 'end_date', e.target.value)}
                                        size="small" InputLabelProps={{ shrink: true }} />
                                    <IconButton size="small" onClick={() => removeTerm(i)}>
                                        <DeleteIcon fontSize="small" />
                                    </IconButton>
                                </Box>
                                <TextField
                                    label="Notes (optional)"
                                    value={term.notes}
                                    onChange={e => updateTerm(i, 'notes', e.target.value)}
                                    size="small"
                                    fullWidth
                                    multiline
                                    minRows={1}
                                    placeholder="Any notes about this term"
                                    sx={{ pl: 0 }}
                                />
                            </Box>
                        ))}
                    </Box>
                )}

                {/* ── Step 2: Term Breaks ──────────────────────────────────── */}
                {step === 2 && (
                    <Box>
                        <Box sx={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', mb: 0.5 }}>
                            <Typography variant="subtitle2">Term Breaks</Typography>
                            <Button size="small" startIcon={<AddIcon />} onClick={addTermBreak}>Add Break</Button>
                        </Box>
                        <Typography variant="caption" sx={{ color: 'text.secondary', display: 'block', mb: 2 }}>
                            Named holiday periods between terms (Christmas, Easter, half-terms etc.)
                        </Typography>
                        {termBreaks.length === 0 && (
                            <Typography variant="body2" sx={{ color: 'text.secondary', fontStyle: 'italic' }}>
                                No term breaks defined. Click "Add Break" to add one, or skip this step.
                            </Typography>
                        )}
                        {termBreaks.map((tb, i) => (
                            <Box key={i} sx={{ display: 'flex', gap: 1.5, mb: 1.5, alignItems: 'center' }}>
                                <TextField label="Break Name" value={tb.name}
                                    onChange={e => updateTermBreak(i, 'name', e.target.value)}
                                    size="small" sx={{ width: 180 }}
                                    placeholder="e.g. Christmas Break" />
                                <TextField label="Start Date" type="date" value={tb.start_date}
                                    onChange={e => updateTermBreak(i, 'start_date', e.target.value)}
                                    size="small" InputLabelProps={{ shrink: true }} />
                                <TextField label="End Date" type="date" value={tb.end_date}
                                    onChange={e => updateTermBreak(i, 'end_date', e.target.value)}
                                    size="small" InputLabelProps={{ shrink: true }} />
                                <IconButton size="small" onClick={() => removeTermBreak(i)}>
                                    <DeleteIcon fontSize="small" />
                                </IconButton>
                            </Box>
                        ))}
                    </Box>
                )}

                {/* ── Step 3: Daily Periods ────────────────────────────────── */}
                {step === 3 && (
                    <Box>
                        <Box sx={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', mb: 1.5 }}>
                            <Typography variant="subtitle2">Daily Period Schedule</Typography>
                            <Button size="small" startIcon={<AddIcon />} onClick={addPeriod}>Add Period</Button>
                        </Box>
                        {periods.map((p, i) => (
                            <Box key={i} sx={{ display: 'flex', gap: 1.5, mb: 1.5, alignItems: 'center' }}>
                                <TextField label="Code" value={p.code}
                                    onChange={e => updatePeriod(i, 'code', e.target.value)}
                                    size="small" sx={{ width: 80 }} />
                                <TextField label="Name" value={p.name}
                                    onChange={e => updatePeriod(i, 'name', e.target.value)}
                                    size="small" sx={{ width: 140 }} />
                                <TextField label="Start" type="time" value={p.start_time}
                                    onChange={e => updatePeriod(i, 'start_time', e.target.value)}
                                    size="small" InputLabelProps={{ shrink: true }} />
                                <TextField label="End" type="time" value={p.end_time}
                                    onChange={e => updatePeriod(i, 'end_time', e.target.value)}
                                    size="small" InputLabelProps={{ shrink: true }} />
                                <FormControl size="small" sx={{ width: 150 }}>
                                    <InputLabel>Type</InputLabel>
                                    <Select label="Type" value={p.period_type}
                                        onChange={e => updatePeriod(i, 'period_type', e.target.value)}>
                                        <MenuItem value="lesson">Lesson</MenuItem>
                                        <MenuItem value="break">Break</MenuItem>
                                        <MenuItem value="registration">Registration</MenuItem>
                                        <MenuItem value="offtimetable">Off-timetable</MenuItem>
                                    </Select>
                                </FormControl>
                                <IconButton size="small" onClick={() => removePeriod(i)}>
                                    <DeleteIcon fontSize="small" />
                                </IconButton>
                            </Box>
                        ))}
                    </Box>
                )}

                {/* ── Step 4: Week Cycles ──────────────────────────────────── */}
                {step === 4 && (
                    <Box>
                        <Box sx={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center', mb: 0.5 }}>
                            <Typography variant="subtitle2">Week A/B Cycles</Typography>
                            <Tooltip title="Reset to alternating A/B (default)">
                                <Button size="small" onClick={resetWeekCycles}>Reset Defaults</Button>
                            </Tooltip>
                        </Box>
                        <Typography variant="caption" sx={{ color: 'text.secondary', display: 'block', mb: 2 }}>
                            Default: Week 1 = A, Week 2 = B, alternating within each term. Click a chip to toggle.
                        </Typography>
                        {weekEntries.length === 0 && (
                            <Typography variant="body2" sx={{ color: 'text.secondary', fontStyle: 'italic' }}>
                                No weeks found — check your term dates.
                            </Typography>
                        )}
                        {Object.entries(weeksByTerm).map(([termName, weeks]) => (
                            <Box key={termName} sx={{ mb: 2 }}>
                                <Typography variant="caption" sx={{ fontWeight: 600, color: 'text.secondary', mb: 0.75, display: 'block' }}>
                                    {termName}
                                </Typography>
                                <Box sx={{ display: 'flex', flexWrap: 'wrap', gap: 0.75 }}>
                                    {weeks.map((w) => {
                                        const globalIdx = weekEntries.findIndex(
                                            e => e.termNumber === w.termNumber && e.weekNumber === w.weekNumber
                                        );
                                        return (
                                            <Tooltip
                                                key={`${w.termNumber}-${w.weekNumber}`}
                                                title={`w/c ${w.startDate}`}
                                            >
                                                <Chip
                                                    label={`W${w.weekNumber} ${w.cycle}`}
                                                    size="small"
                                                    color={w.cycle === 'A' ? 'primary' : 'secondary'}
                                                    variant="outlined"
                                                    onClick={() => toggleWeekCycle(globalIdx)}
                                                    sx={{ cursor: 'pointer', fontWeight: 600, minWidth: 56 }}
                                                />
                                            </Tooltip>
                                        );
                                    })}
                                </Box>
                            </Box>
                        ))}
                    </Box>
                )}
            </DialogContent>

            <DialogActions sx={{ px: 3, pb: 2 }}>
                <Button onClick={handleClose} disabled={saving}>Cancel</Button>
                {step > 0 && (
                    <Button onClick={() => setStep(s => s - 1)} disabled={saving}>Back</Button>
                )}
                {step === 0 && (
                    <Button onClick={handleSaveSchoolInfo} variant="contained" disabled={saving}>
                        {saving ? <CircularProgress size={18} /> : 'Save & Continue'}
                    </Button>
                )}
                {step === 1 && (
                    <Button onClick={() => setStep(2)} variant="outlined" disabled={saving}>
                        Next: Term Breaks
                    </Button>
                )}
                {step === 2 && (
                    <Button onClick={() => setStep(3)} variant="outlined" disabled={saving}>
                        Next: Daily Periods
                    </Button>
                )}
                {step === 3 && (
                    <Button onClick={() => goToStep(4)} variant="outlined" disabled={saving}>
                        Next: Week Cycles
                    </Button>
                )}
                {step === 4 && (
                    <Button onClick={handleSaveCalendar} variant="contained" disabled={saving}>
                        {saving ? <CircularProgress size={18} /> : 'Save School Calendar'}
                    </Button>
                )}
            </DialogActions>
        </Dialog>
    );
}
