auditly/contexts/UserOrganizationsContext.tsx

import React, { createContext, useContext, useEffect, useState } from 'react';
import { useAuth } from './AuthContext';
import { isFirebaseConfigured } from '../services/firebase';
import { API_URL } from '../constants';
import { demoStorage } from '../services/demoStorage';

interface UserOrganization {
    orgId: string;
    name: string;
    role: 'owner' | 'admin' | 'employee';
    onboardingCompleted: boolean;
    joinedAt: number;
}

interface UserOrganizationsContextType {
    organizations: UserOrganization[];
    selectedOrgId: string | null;
    loading: boolean;
    selectOrganization: (orgId: string) => void;
    createOrganization: (name: string) => Promise<{ orgId: string; requiresSubscription?: boolean }>;
    joinOrganization: (inviteCode: string) => Promise<string>;
    refreshOrganizations: () => Promise<void>;
    createCheckoutSession: (orgId: string, userEmail: string) => Promise<{ sessionUrl: string; sessionId: string }>;
    getSubscriptionStatus: (orgId: string) => Promise<any>;
}

const UserOrganizationsContext = createContext<UserOrganizationsContextType | undefined>(undefined);

export const UserOrganizationsProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
    const { user } = useAuth();
    const [organizations, setOrganizations] = useState<UserOrganization[]>([]);
    const [selectedOrgId, setSelectedOrgId] = useState<string | null>(null);
    const [loading, setLoading] = useState(true);

    // Load user's organizations
    const loadOrganizations = async () => {
        if (!user) {
            setOrganizations([]);
            setLoading(false);
            return;
        }

        try {
            if (!isFirebaseConfigured) {
                // Demo mode - fetch from server API
                const response = await fetch(`${API_URL}/api/user/${user.uid}/organizations`);
                if (response.ok) {
                    const data = await response.json();
                    setOrganizations(data.organizations || []);
                } else {
                    console.error('Failed to load organizations:', response.status);
                    setOrganizations([]);
                }
            } else {
                // Firebase mode - fetch from Cloud Functions
                const response = await fetch(`${API_URL}/getUserOrganizations?userId=${user.uid}`);
                if (response.ok) {
                    const data = await response.json();
                    setOrganizations(data.organizations || []);
                } else {
                    console.error('Failed to load organizations:', response.status);
                    setOrganizations([]);
                }
            }
        } catch (error) {
            console.error('Failed to load organizations:', error);
            setOrganizations([]);
        } finally {
            setLoading(false);
        }
    };

    // Initialize selected org from localStorage (persistent across sessions)
    useEffect(() => {
        const savedOrgId = localStorage.getItem('auditly_selected_org');
        if (savedOrgId) {
            setSelectedOrgId(savedOrgId);
        }
    }, []);

    // Load organizations when user changes
    useEffect(() => {
        loadOrganizations();
    }, [user]);

    // Listen for organization updates (e.g., onboarding completion)
    useEffect(() => {
        const handleOrgUpdate = (event: CustomEvent) => {
            const { orgId, onboardingCompleted } = event.detail;
            console.log('UserOrganizationsContext received org update:', { orgId, onboardingCompleted });

            if (onboardingCompleted && orgId) {
                // Update the specific organization in the list to reflect onboarding completion
                setOrganizations(prev => {
                    const updated = prev.map(org =>
                        org.orgId === orgId
                            ? { ...org, onboardingCompleted: true }
                            : org
                    );
                    console.log('Updated organizations after onboarding completion:', updated);
                    return updated;
                });
            }
        };

        window.addEventListener('organizationUpdated', handleOrgUpdate as EventListener);

        return () => {
            window.removeEventListener('organizationUpdated', handleOrgUpdate as EventListener);
        };
    }, []);

    const selectOrganization = (orgId: string) => {
        console.log('Switching to organization:', orgId);

        // Clear any cached data when switching organizations for security
        sessionStorage.removeItem('auditly_cached_employees');
        sessionStorage.removeItem('auditly_cached_submissions');
        sessionStorage.removeItem('auditly_cached_reports');

        setSelectedOrgId(orgId);
        localStorage.setItem('auditly_selected_org', orgId);

        // Dispatch event to notify other contexts about the org switch
        window.dispatchEvent(new CustomEvent('organizationChanged', {
            detail: { newOrgId: orgId }
        }));
    };

    const createOrganization = async (name: string): Promise<{ orgId: string; requiresSubscription?: boolean }> => {
        if (!user) throw new Error('User not authenticated');

        try {
            let newOrg: UserOrganization;
            let requiresSubscription = false;

            if (!isFirebaseConfigured) {
                // Demo mode - use server API
                const response = await fetch(`${API_URL}/api/organizations`, {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ name, userId: user.uid })
                });

                if (!response.ok) {
                    throw new Error(`Failed to create organization: ${response.status}`);
                }

                const data = await response.json();
                newOrg = {
                    orgId: data.orgId,
                    name: data.name,
                    role: data.role,
                    onboardingCompleted: data.onboardingCompleted,
                    joinedAt: data.joinedAt
                };

                setOrganizations(prev => [...prev, newOrg]);
            } else {
                // Firebase mode - use Cloud Function
                const response = await fetch(`${API_URL}/createOrganization`, {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ name, userId: user.uid })
                });

                if (!response.ok) {
                    throw new Error(`Failed to create organization: ${response.status}`);
                }

                const data = await response.json();
                newOrg = {
                    orgId: data.orgId,
                    name: data.name,
                    role: data.role,
                    onboardingCompleted: data.onboardingCompleted,
                    joinedAt: data.joinedAt
                };

                requiresSubscription = data.requiresSubscription || false;
                setOrganizations(prev => [...prev, newOrg]);
            }

            return { orgId: newOrg.orgId, requiresSubscription };
        } catch (error) {
            console.error('Failed to create organization:', error);
            throw error;
        }
    };

    const joinOrganization = async (inviteCode: string): Promise<string> => {
        if (!user) throw new Error('User not authenticated');

        try {
            if (!isFirebaseConfigured) {
                // Demo mode - use server API to get and consume invite
                const inviteStatusRes = await fetch(`/api/invitations/${inviteCode}`);
                if (!inviteStatusRes.ok) {
                    throw new Error('Invalid or expired invite code');
                }

                const inviteData = await inviteStatusRes.json();
                if (inviteData.used) {
                    throw new Error('Invite code has already been used');
                }

                // Consume the invite
                const consumeRes = await fetch(`/api/invitations/${inviteCode}/consume`, {
                    method: 'POST'
                });
                if (!consumeRes.ok) {
                    throw new Error('Failed to consume invite');
                }

                const consumedData = await consumeRes.json();
                const orgId = consumedData.orgId;

                // Get organization data (this might be from localStorage for demo mode)
                const orgData = demoStorage.getOrganization(orgId);
                if (!orgData) {
                    throw new Error('Organization not found');
                }

                const userOrg: UserOrganization = {
                    orgId: orgId,
                    name: orgData.name,
                    role: 'employee',
                    onboardingCompleted: orgData.onboardingCompleted || false,
                    joinedAt: Date.now()
                };

                setOrganizations(prev => [...prev, userOrg]);
                return orgId;
            } else {
                // Firebase mode - use Cloud Function
                const response = await fetch(`${API_URL}/joinOrganization`, {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ userId: user.uid, inviteCode })
                });

                if (!response.ok) {
                    const errorData = await response.json();
                    throw new Error(errorData.error || 'Failed to join organization');
                }

                const data = await response.json();
                const userOrg: UserOrganization = {
                    orgId: data.orgId,
                    name: data.name,
                    role: data.role,
                    onboardingCompleted: data.onboardingCompleted,
                    joinedAt: data.joinedAt
                };

                setOrganizations(prev => [...prev, userOrg]);
                return data.orgId;
            }
        } catch (error) {
            console.error('Failed to join organization:', error);
            throw error;
        }
    };

    const refreshOrganizations = async () => {
        setLoading(true);
        await loadOrganizations();
    };

    const createCheckoutSession = async (orgId: string, userEmail: string): Promise<{ sessionUrl: string; sessionId: string }> => {
        if (!user) throw new Error('User not authenticated');

        try {
            const response = await fetch(`${API_URL}/createCheckoutSession`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    orgId,
                    userId: user.uid,
                    userEmail
                })
            });

            if (!response.ok) {
                const errorData = await response.json();
                throw new Error(errorData.error || 'Failed to create checkout session');
            }

            const data = await response.json();
            return {
                sessionUrl: data.sessionUrl,
                sessionId: data.sessionId
            };
        } catch (error) {
            console.error('Failed to create checkout session:', error);
            throw error;
        }
    };

    const getSubscriptionStatus = async (orgId: string) => {
        try {
            const response = await fetch(`${API_URL}/getSubscriptionStatus?orgId=${orgId}`);

            if (!response.ok) {
                throw new Error('Failed to get subscription status');
            }

            const data = await response.json();
            return data;
        } catch (error) {
            console.error('Failed to get subscription status:', error);
            throw error;
        }
    };

    return (
        <UserOrganizationsContext.Provider value={{
            organizations,
            selectedOrgId,
            loading,
            selectOrganization,
            createOrganization,
            joinOrganization,
            refreshOrganizations,
            createCheckoutSession,
            getSubscriptionStatus
        }}>
            {children}
        </UserOrganizationsContext.Provider>
    );
};

export const useUserOrganizations = () => {
    const context = useContext(UserOrganizationsContext);
    if (!context) {
        throw new Error('useUserOrganizations must be used within UserOrganizationsProvider');
    }
    return context;
};