frontend/src/context/GeneratorWorkflowContext.jsx

import React, {
    createContext,
    useCallback,
    useContext,
    useEffect,
    useMemo,
    useRef,
    useState,
} from 'react';
import { apiFetch } from '@/lib/api';
import { mergeSequenceIntoContacts } from '@/lib/mergeSequenceIntoContacts';

const GeneratorWorkflowContext = createContext(null);

export function GeneratorWorkflowProvider({ children }) {
    const [step, setStep] = useState(1);
    const [uploadedFile, setUploadedFile] = useState(null);
    const [selectedProducts, setSelectedProducts] = useState([]);
    const [prompts, setPrompts] = useState({});
    const [linkedinPrompts, setLinkedinPrompts] = useState({});
    const [includeLinkedin, setIncludeLinkedin] = useState(false);
    const [isGenerating, setIsGenerating] = useState(false);
    const [generationComplete, setGenerationComplete] = useState(false);
    const [generationRunId, setGenerationRunId] = useState(0);
    const [genPhase, setGenPhase] = useState('email');
    const [contacts, setContacts] = useState([]);
    const [sequences, setSequences] = useState([]);
    const [progress, setProgress] = useState(0);
    const [streamComplete, setStreamComplete] = useState(false);
    const [reconnectKey, setReconnectKey] = useState(0);

    const prevRunIdRef = useRef(null);
    const eventSourceRef = useRef(null);

    const resetWorkflow = useCallback(() => {
        if (eventSourceRef.current) {
            eventSourceRef.current.close();
            eventSourceRef.current = null;
        }
        setStep(1);
        setUploadedFile(null);
        setSelectedProducts([]);
        setPrompts({});
        setLinkedinPrompts({});
        setIncludeLinkedin(false);
        setIsGenerating(false);
        setGenerationComplete(false);
        setGenerationRunId(0);
        setGenPhase('email');
        setContacts([]);
        setSequences([]);
        setProgress(0);
        setStreamComplete(false);
        prevRunIdRef.current = null;
    }, []);

    const beginGeneration = useCallback((runIdIncrement = true) => {
        setStreamComplete(false);
        setGenerationComplete(false);
        setContacts([]);
        setSequences([]);
        setProgress(0);
        setGenPhase('email');
        setIsGenerating(true);
        if (runIdIncrement) {
            setGenerationRunId((r) => r + 1);
        }
    }, []);

    const contactCount = uploadedFile?.contactCount || 0;

    useEffect(() => {
        if (!isGenerating || !uploadedFile?.fileId) return;

        const isNewRun = prevRunIdRef.current !== generationRunId;
        if (isNewRun) {
            prevRunIdRef.current = generationRunId;
            setContacts([]);
            setSequences([]);
            setProgress(0);
            setStreamComplete(false);
        }

        const reset = isNewRun ? 1 : 0;
        const url = `/api/generate-sequences?file_id=${encodeURIComponent(uploadedFile.fileId)}&reset=${reset}`;
        const eventSource = new EventSource(url);
        eventSourceRef.current = eventSource;

        eventSource.onmessage = (event) => {
            try {
                const data = JSON.parse(event.data);
                if (data.type === 'sequence') {
                    const seq = data.sequence;
                    setSequences((prev) => {
                        const sig = `${seq.id}-${seq.emailNumber}-${seq.channel || 'email'}-${seq.stepOrder ?? ''}`;
                        if (
                            prev.some(
                                (s) =>
                                    `${s.id}-${s.emailNumber}-${s.channel || 'email'}-${s.stepOrder ?? ''}` === sig
                            )
                        ) {
                            return prev;
                        }
                        return [...prev, seq];
                    });
                    setContacts((prev) => mergeSequenceIntoContacts(prev, seq));
                } else if (data.type === 'progress') {
                    if (typeof data.progress === 'number') setProgress(data.progress);
                } else if (data.type === 'phase') {
                    if (data.phase === 'linkedin') setGenPhase('linkedin');
                } else if (data.type === 'complete') {
                    setStreamComplete(true);
                    setIsGenerating(false);
                    setGenerationComplete(true);
                    eventSource.close();
                    eventSourceRef.current = null;
                } else if (data.type === 'error') {
                    console.error('Generation error:', data.error);
                    alert('Error generating sequences: ' + data.error);
                    eventSource.close();
                    eventSourceRef.current = null;
                    setIsGenerating(false);
                }
            } catch (err) {
                console.error('Error parsing SSE data:', err);
            }
        };

        eventSource.onerror = () => {
            eventSource.close();
            eventSourceRef.current = null;
            if (!streamComplete) setReconnectKey((k) => k + 1);
        };

        return () => {
            eventSource.close();
            if (eventSourceRef.current === eventSource) eventSourceRef.current = null;
        };
    }, [isGenerating, uploadedFile?.fileId, generationRunId]);

    useEffect(() => {
        if (!isGenerating || !uploadedFile?.fileId || reconnectKey === 0) return;
        let cancelled = false;
        (async () => {
            try {
                const [statusRes, seqRes] = await Promise.all([
                    apiFetch(`/api/generation-status?file_id=${encodeURIComponent(uploadedFile.fileId)}`),
                    apiFetch(`/api/sequences?file_id=${encodeURIComponent(uploadedFile.fileId)}`),
                ]);
                if (cancelled) return;
                if (statusRes.ok && seqRes.ok) {
                    const status = await statusRes.json();
                    const { sequences: list } = await seqRes.json();
                    if (status.is_complete) {
                        setStreamComplete(true);
                        setIsGenerating(false);
                        setGenerationComplete(true);
                    }
                    if (list?.length > 0) {
                        const byContact = new Map();
                        list.forEach((seq) => {
                            const key = seq.email;
                            const ch = seq.channel || 'email';
                            if (!byContact.has(key)) {
                                byContact.set(key, {
                                    id: seq.id,
                                    firstName: seq.firstName,
                                    lastName: seq.lastName,
                                    email: seq.email,
                                    company: seq.company,
                                    title: seq.title,
                                    product: seq.product,
                                    emails: [],
                                    linkedin: [],
                                });
                            }
                            const row = {
                                emailNumber: seq.emailNumber,
                                subject: seq.subject,
                                emailContent: seq.emailContent,
                                channel: ch,
                                stepOrder: seq.stepOrder,
                            };
                            const dedupe = (e) =>
                                e.stepOrder != null
                                    ? `ord-${e.stepOrder}-${ch}`
                                    : `${ch}-${e.emailNumber}`;
                            const rowKey = dedupe(row);
                            const c = byContact.get(key);
                            if (ch === 'linkedin') {
                                if (!c.linkedin.some((e) => dedupe(e) === rowKey)) {
                                    c.linkedin.push(row);
                                }
                            } else {
                                if (!c.emails.some((e) => dedupe(e) === rowKey)) {
                                    c.emails.push(row);
                                }
                            }
                        });
                        const arr = [...byContact.values()];
                        arr.sort((a, b) => (a.id || 0) - (b.id || 0));
                        setSequences(list);
                        setContacts(arr);
                        if (status.total_contacts > 0) {
                            const exp =
                                status.campaign_sequence_actions && status.total_contacts > 0
                                    ? status.total_contacts * status.campaign_sequence_actions
                                    : status.has_linkedin_prompts
                                      ? status.total_contacts * 2
                                      : status.total_contacts;
                            setProgress(
                                Math.min(100, ((status.completed_count || 0) / exp) * 100)
                            );
                        }
                    }
                }
            } catch (e) {
                if (!cancelled) console.error('Reconnect fetch error:', e);
            }
        })();
        return () => {
            cancelled = true;
        };
    }, [reconnectKey, isGenerating, uploadedFile?.fileId, contactCount]);

    useEffect(() => {
        if (!isGenerating || !uploadedFile?.fileId) return;
        const onVisible = () => {
            if (document.visibilityState === 'visible') setReconnectKey((k) => k + 1);
        };
        document.addEventListener('visibilitychange', onVisible);
        return () => document.removeEventListener('visibilitychange', onVisible);
    }, [isGenerating, uploadedFile?.fileId]);

    const value = useMemo(
        () => ({
            step,
            setStep,
            uploadedFile,
            setUploadedFile,
            selectedProducts,
            setSelectedProducts,
            prompts,
            setPrompts,
            linkedinPrompts,
            setLinkedinPrompts,
            includeLinkedin,
            setIncludeLinkedin,
            isGenerating,
            generationComplete,
            generationRunId,
            beginGeneration,
            setIsGenerating,
            setGenerationComplete,
            resetWorkflow,
            contacts,
            sequences,
            progress,
            genPhase,
            streamComplete,
            contactCount,
        }),
        [
            step,
            uploadedFile,
            selectedProducts,
            prompts,
            linkedinPrompts,
            includeLinkedin,
            isGenerating,
            generationComplete,
            generationRunId,
            beginGeneration,
            resetWorkflow,
            contacts,
            sequences,
            progress,
            genPhase,
            streamComplete,
            contactCount,
        ]
    );

    return (
        <GeneratorWorkflowContext.Provider value={value}>{children}</GeneratorWorkflowContext.Provider>
    );
}

export function useGeneratorWorkflow() {
    const ctx = useContext(GeneratorWorkflowContext);
    if (!ctx) {
        throw new Error('useGeneratorWorkflow must be used within GeneratorWorkflowProvider');
    }
    return ctx;
}