import React, { useState, useCallback, useRef } from 'react';
import {
    Dialog, DialogTitle, DialogContent, DialogActions,
    Button, Stepper, Step, StepLabel, Box, TextField,
    Typography, CircularProgress, Alert, List, ListItem,
    ListItemButton, ListItemText, Divider, InputAdornment,
    Chip,
} from '@mui/material';
import { Search as SearchIcon, School as SchoolIcon } from '@mui/icons-material';
import { useAuth } from '../../../../../../contexts/AuthContext';
import { SchoolCalendarWizard, SchoolInfo } from './SchoolCalendarWizard';

// ─── Types ────────────────────────────────────────────────────────────────────

interface GaisSchool {
    urn: string;
    name: string;
    status: string | null;
    phase: string | null;
    type: string | null;
    street: string | null;
    locality: string | null;
    town: string | null;
    county: string | null;
    postcode: string | null;
    website: string | null;
    telephone: string | null;
    head_title: string | null;
    head_first_name: string | null;
    head_last_name: string | null;
    la_name: string | null;
    number_of_pupils: number | null;
    region: string | null;
}

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

const STEPS = ['Find your school', 'Confirm details', 'Set up calendar'];

// ─── Helpers ──────────────────────────────────────────────────────────────────

function buildAddress(g: GaisSchool): Record<string, string> {
    const addr: Record<string, string> = {};
    if (g.street) addr.street = g.street;
    if (g.locality) addr.locality = g.locality;
    if (g.town) addr.town = g.town;
    if (g.county) addr.county = g.county;
    if (g.postcode) addr.postcode = g.postcode;
    return addr;
}

function headteacherString(g: GaisSchool): string {
    return [g.head_title, g.head_first_name, g.head_last_name]
        .filter(Boolean)
        .join(' ');
}

// ─── Component ────────────────────────────────────────────────────────────────

export function SchoolOnboardingWizard({ open, onClose, onComplete, apiBase }: Props) {
    const { accessToken } = useAuth();

    const [step, setStep] = useState(0);
    const [searching, setSearching] = useState(false);
    const [registering, setRegistering] = useState(false);
    const [error, setError] = useState<string | null>(null);

    // Step 0 — search
    const [query, setQuery] = useState('');
    const [results, setResults] = useState<GaisSchool[]>([]);
    const [selected, setSelected] = useState<GaisSchool | null>(null);
    const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);

    // Step 1 — confirm/edit
    const [editName, setEditName] = useState('');
    const [editWebsite, setEditWebsite] = useState('');
    const [editHeadteacher, setEditHeadteacher] = useState('');

    // After registration
    const [registeredSchoolInfo, setRegisteredSchoolInfo] = useState<SchoolInfo | null>(null);
    const [calendarOpen, setCalendarOpen] = useState(false);

    const reset = () => {
        setStep(0);
        setQuery('');
        setResults([]);
        setSelected(null);
        setEditName('');
        setEditWebsite('');
        setEditHeadteacher('');
        setRegisteredSchoolInfo(null);
        setCalendarOpen(false);
        setError(null);
    };

    const handleClose = () => { reset(); onClose(); };

    // ── Step 0: search ────────────────────────────────────────────────────────

    const doSearch = useCallback(async (q: string) => {
        if (q.trim().length < 2) { setResults([]); return; }
        setSearching(true);
        setError(null);
        try {
            const res = await fetch(
                `${apiBase}/school/search?q=${encodeURIComponent(q)}&limit=25&status=Open`,
                { headers: { Authorization: `Bearer ${accessToken}` } }
            );
            const data = await res.json();
            if (data.status === 'ok') setResults(data.schools || []);
            else setError(data.message || 'Search failed');
        } catch {
            setError('Search request failed');
        } finally {
            setSearching(false);
        }
    }, [apiBase, accessToken]);

    const handleQueryChange = (e: React.ChangeEvent<HTMLInputElement>) => {
        const q = e.target.value;
        setQuery(q);
        if (debounceRef.current) clearTimeout(debounceRef.current);
        debounceRef.current = setTimeout(() => doSearch(q), 350);
    };

    const handleSelectSchool = (school: GaisSchool) => {
        setSelected(school);
        setEditName(school.name);
        setEditWebsite(school.website || '');
        setEditHeadteacher(headteacherString(school));
        setStep(1);
    };

    // ── Step 1 → Step 2: register school ─────────────────────────────────────

    const handleRegister = async () => {
        if (!selected) return;
        setRegistering(true);
        setError(null);
        try {
            const res = await fetch(`${apiBase}/school/register`, {
                method: 'POST',
                headers: {
                    Authorization: `Bearer ${accessToken}`,
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    urn: selected.urn,
                    name: editName,
                    address: buildAddress(selected),
                    website: editWebsite,
                    headteacher: editHeadteacher,
                }),
            });
            const data = await res.json();
            if (data.status === 'ok' || data.status === 'already_exists') {
                const info: SchoolInfo = {
                    name: editName,
                    urn: selected.urn,
                    website: editWebsite,
                    address: buildAddress(selected),
                    headteacher: editHeadteacher,
                    term_dates_url: '',
                    staff_list_url: '',
                };
                setRegisteredSchoolInfo(info);
                setStep(2);
            } else {
                setError(data.message || 'Registration failed');
            }
        } catch {
            setError('Registration request failed');
        } finally {
            setRegistering(false);
        }
    };

    // ── Step 2: calendar wizard completes → done ──────────────────────────────

    const handleCalendarComplete = () => {
        setCalendarOpen(false);
        reset();
        onComplete();
    };

    // ── Render ────────────────────────────────────────────────────────────────

    return (
        <>
            <Dialog open={open && step < 2} onClose={handleClose} maxWidth="sm" fullWidth>
                <DialogTitle sx={{ pb: 1 }}>
                    <Box sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
                        <SchoolIcon fontSize="small" />
                        Set up your school
                    </Box>
                </DialogTitle>

                <DialogContent>
                    <Stepper activeStep={step} sx={{ mb: 3, mt: 0.5 }}>
                        {STEPS.map(label => (
                            <Step key={label}>
                                <StepLabel>{label}</StepLabel>
                            </Step>
                        ))}
                    </Stepper>

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

                    {step === 0 && (
                        <Box>
                            <TextField
                                fullWidth
                                label="Search by school name, URN, or postcode"
                                value={query}
                                onChange={handleQueryChange}
                                size="small"
                                InputProps={{
                                    startAdornment: (
                                        <InputAdornment position="start">
                                            {searching ? <CircularProgress size={16} /> : <SearchIcon fontSize="small" />}
                                        </InputAdornment>
                                    ),
                                }}
                                sx={{ mb: 1 }}
                            />
                            {results.length > 0 && (
                                <List dense sx={{ maxHeight: 360, overflow: 'auto', border: '1px solid', borderColor: 'divider', borderRadius: 1 }}>
                                    {results.map((school, idx) => (
                                        <React.Fragment key={school.urn}>
                                            {idx > 0 && <Divider component="li" />}
                                            <ListItem disablePadding>
                                                <ListItemButton onClick={() => handleSelectSchool(school)}>
                                                    <ListItemText
                                                        primary={school.name}
                                                        secondary={
                                                            <Box component="span" sx={{ display: 'flex', flexWrap: 'wrap', gap: 0.5, mt: 0.25 }}>
                                                                <Typography component="span" variant="caption" color="text.secondary">
                                                                    {[school.town, school.county, school.postcode].filter(Boolean).join(', ')}
                                                                </Typography>
                                                                {school.phase && (
                                                                    <Chip label={school.phase} size="small" variant="outlined" sx={{ height: 16, fontSize: '0.65rem' }} />
                                                                )}
                                                            </Box>
                                                        }
                                                        primaryTypographyProps={{ variant: 'body2', fontWeight: 500 }}
                                                    />
                                                </ListItemButton>
                                            </ListItem>
                                        </React.Fragment>
                                    ))}
                                </List>
                            )}
                            {!searching && query.length >= 2 && results.length === 0 && (
                                <Typography variant="caption" color="text.secondary">
                                    No open schools found. Try a different name or postcode.
                                </Typography>
                            )}
                        </Box>
                    )}

                    {step === 1 && selected && (
                        <Box sx={{ display: 'flex', flexDirection: 'column', gap: 2 }}>
                            <Typography variant="body2" color="text.secondary">
                                Review and edit your school details before registering.
                            </Typography>
                            <TextField
                                fullWidth
                                label="School name"
                                value={editName}
                                onChange={e => setEditName(e.target.value)}
                                size="small"
                            />
                            <TextField
                                fullWidth
                                label="Website"
                                value={editWebsite}
                                onChange={e => setEditWebsite(e.target.value)}
                                size="small"
                                placeholder="https://"
                            />
                            <TextField
                                fullWidth
                                label="Headteacher"
                                value={editHeadteacher}
                                onChange={e => setEditHeadteacher(e.target.value)}
                                size="small"
                            />
                            <Box sx={{ p: 1.5, borderRadius: 1, bgcolor: 'action.hover', fontSize: '0.78rem' }}>
                                <Typography variant="caption" display="block" color="text.secondary" sx={{ mb: 0.5 }}>Address from Edubase</Typography>
                                {Object.values(buildAddress(selected)).filter(Boolean).join(', ')}
                            </Box>
                            <Typography variant="caption" color="text.secondary">
                                URN: {selected.urn} · {selected.phase} · {selected.type}
                            </Typography>
                        </Box>
                    )}
                </DialogContent>

                <DialogActions>
                    {step === 0 && (
                        <Button onClick={handleClose} size="small">Cancel</Button>
                    )}
                    {step === 1 && (
                        <>
                            <Button onClick={() => setStep(0)} size="small">Back</Button>
                            <Button onClick={handleClose} size="small">Cancel</Button>
                            <Button
                                variant="contained"
                                size="small"
                                onClick={handleRegister}
                                disabled={registering || !editName.trim()}
                                startIcon={registering ? <CircularProgress size={14} /> : undefined}
                            >
                                {registering ? 'Registering…' : 'Register school'}
                            </Button>
                        </>
                    )}
                </DialogActions>
            </Dialog>

            {/* Step 2: calendar setup opens as a separate wizard */}
            {registeredSchoolInfo && (
                <SchoolCalendarWizard
                    open={step === 2}
                    onClose={() => { reset(); onClose(); }}
                    onComplete={handleCalendarComplete}
                    apiBase={apiBase}
                    schoolInfo={registeredSchoolInfo}
                />
            )}
        </>
    );
}
