gymflow/src/context/DataContext.tsx

import React, { createContext, useContext, useState, useEffect } from 'react';
import { WorkoutSession, WorkoutPlan, WorkoutSet } from '../types';
import { useAuth } from './AuthContext';
import {
    getSessions,
    getPlans,
    getActiveSession,
    saveSession,
    deleteSession,
    updateActiveSession,
    deleteActiveSession,
    deleteSetFromActiveSession,
    updateSetInActiveSession
} from '../services/storage';
import { getCurrentUserProfile, getMe } from '../services/auth';
import { generateId } from '../utils/uuid';
import { logWeight } from '../services/weight';
import { useNavigate } from 'react-router-dom';

interface DataContextType {
    sessions: WorkoutSession[];
    plans: WorkoutPlan[];
    activeSession: WorkoutSession | null;
    activePlan: WorkoutPlan | null;
    startSession: (plan?: WorkoutPlan, startWeight?: number) => Promise<void>;
    endSession: () => Promise<void>;
    quitSession: () => Promise<void>;
    addSet: (set: WorkoutSet) => void;
    removeSet: (setId: string) => Promise<void>;
    updateSet: (updatedSet: WorkoutSet) => Promise<void>;
    updateSession: (updatedSession: WorkoutSession) => void;
    deleteSessionById: (sessionId: string) => void;
    refreshData: () => Promise<void>;
}

const DataContext = createContext<DataContextType | undefined>(undefined);

export const DataProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
    const { currentUser, updateUser } = useAuth();
    const navigate = useNavigate();

    const [sessions, setSessions] = useState<WorkoutSession[]>([]);
    const [plans, setPlans] = useState<WorkoutPlan[]>([]);
    const [activeSession, setActiveSession] = useState<WorkoutSession | null>(null);
    const [activePlan, setActivePlan] = useState<WorkoutPlan | null>(null);

    const refreshData = async () => {
        if (currentUser) {
            const s = await getSessions(currentUser.id);
            setSessions(s);
            const p = await getPlans(currentUser.id);
            setPlans(p);
        } else {
            setSessions([]);
            setPlans([]);
        }
    };

    useEffect(() => {
        refreshData();
    }, [currentUser]);

    // Restore active session
    useEffect(() => {
        const restoreActive = async () => {
            if (currentUser) {
                const session = await getActiveSession(currentUser.id);
                if (session) {
                    setActiveSession(session);
                    if (session.planId) {
                        // Ensure plans are loaded or fetch specifically
                        const currentPlans = plans.length > 0 ? plans : await getPlans(currentUser.id);
                        const plan = currentPlans.find(p => p.id === session.planId);
                        if (plan) setActivePlan(plan);
                    }
                }
            }
        };
        restoreActive();
    }, [currentUser]); // Dependency logic might need tuning, but this matches App.tsx roughly

    const startSession = async (plan?: WorkoutPlan, startWeight?: number) => {
        if (!currentUser || activeSession) return;

        const profile = getCurrentUserProfile(currentUser.id);
        const currentWeight = startWeight || profile?.weight || 70;

        const newSession: WorkoutSession = {
            id: generateId(),
            startTime: Date.now(),
            type: 'STANDARD',
            userBodyWeight: currentWeight,
            sets: [],
            planId: plan?.id,
            planName: plan?.name
        };

        setActivePlan(plan || null);
        setActiveSession(newSession);
        navigate('/');

        await saveSession(currentUser.id, newSession);

        if (startWeight) {
            await logWeight(startWeight);
        }
    };

    const endSession = async () => {
        if (activeSession && currentUser) {
            const finishedSession = { ...activeSession, endTime: Date.now() };
            await updateActiveSession(currentUser.id, finishedSession);
            setSessions(prev => [finishedSession, ...prev]);
            setActiveSession(null);
            setActivePlan(null);

            const res = await getMe();
            if (res.success && res.user) {
                updateUser(res.user);
            }
        }
    };

    const quitSession = async () => {
        if (currentUser) {
            await deleteActiveSession(currentUser.id);
            setActiveSession(null);
            setActivePlan(null);
        }
    };

    const addSet = (set: WorkoutSet) => {
        if (activeSession) {
            const updatedSession = { ...activeSession, sets: [...activeSession.sets, set] };
            setActiveSession(updatedSession);
            // Context update is optimistic, actual save usually happens in hooks or components? 
            // In App.tsx handleAddSet only updated local state.
            // Wait, useTracker usually handles saving sets via API? 
            // In App.tsx: handleAddSet just set state.
            // useTracker.ts calls onSetAdded, but ALSO calls api to save it?
            // Let's look at useTracker.ts. 
            // handleLogSet in useTracker calls API then onSetAdded.
            // So this state update is mainly for UI sync in App.
        }
    };

    const removeSet = async (setId: string) => {
        if (activeSession && currentUser) {
            await deleteSetFromActiveSession(currentUser.id, setId);
            const updatedSession = {
                ...activeSession,
                sets: activeSession.sets.filter(s => s.id !== setId)
            };
            setActiveSession(updatedSession);
        }
    };

    const updateSet = async (updatedSet: WorkoutSet) => {
        if (activeSession && currentUser) {
            const response = await updateSetInActiveSession(currentUser.id, updatedSet.id, updatedSet);
            const updatedSession = {
                ...activeSession,
                sets: activeSession.sets.map(s => s.id === updatedSet.id ? response : s)
            };
            setActiveSession(updatedSession);
        }
    };

    const updateSession = (updatedSession: WorkoutSession) => {
        if (!currentUser) return;
        saveSession(currentUser.id, updatedSession);
        setSessions(prev => prev.map(s => s.id === updatedSession.id ? updatedSession : s));
    };

    const deleteSessionById = (sessionId: string) => {
        if (!currentUser) return;
        deleteSession(currentUser.id, sessionId);
        setSessions(prev => prev.filter(s => s.id !== sessionId));
    };

    return (
        <DataContext.Provider value={{
            sessions,
            plans,
            activeSession,
            activePlan,
            startSession,
            endSession,
            quitSession,
            addSet,
            removeSet,
            updateSet,
            updateSession,
            deleteSessionById,
            refreshData
        }}>
            {children}
        </DataContext.Provider>
    );
};

export const useData = () => {
    const context = useContext(DataContext);
    if (context === undefined) {
        throw new Error('useData must be used within a DataProvider');
    }
    return context;
};