auditly/pages/OrgSelection.tsx

import React, { useState, useEffect } from 'react';
import { useNavigate, useLocation } from 'react-router-dom';
import { useAuth } from '../contexts/AuthContext';
import { useUserOrganizations } from '../contexts/UserOrganizationsContext';
import { Card, Button } from '../components/UiKit';

const OrgSelection: React.FC = () => {
    const navigate = useNavigate();
    const location = useLocation();
    const { user } = useAuth();
    const {
        organizations,
        loading,
        selectOrganization,
        createOrganization,
        joinOrganization
    } = useUserOrganizations();

    const [showCreateOrg, setShowCreateOrg] = useState(false);
    const [newOrgName, setNewOrgName] = useState('');
    const [inviteCode, setInviteCode] = useState('');
    const [isCreating, setIsCreating] = useState(false);
    const [isJoining, setIsJoining] = useState(false);

    useEffect(() => {
        // Check for invite code in URL
        const hashSearch = location.hash.includes('?') ? location.hash.split('?')[1] : '';
        const searchParams = new URLSearchParams(hashSearch);
        const queryInvite = searchParams.get('invite');
        if (queryInvite) {
            setInviteCode(queryInvite);
        }
    }, [location]);

    // Auto-join organization when invite code is present
    useEffect(() => {
        if (inviteCode && !isJoining && !loading) {
            handleJoinWithInvite();
        }
    }, [inviteCode, isJoining, loading]);

    const handleSelectOrg = (orgId: string) => {
        selectOrganization(orgId);

        // Check if the organization needs onboarding completion
        const selectedOrg = organizations.find(org => org.orgId === orgId);
        if (selectedOrg && !selectedOrg.onboardingCompleted && selectedOrg.role === 'owner') {
            navigate('/onboarding', { replace: true });
        } else {
            navigate('/reports', { replace: true });
        }
    };

    const handleCreateOrg = async () => {
        if (!newOrgName.trim() || isCreating) return;

        setIsCreating(true);
        try {
            const result = await createOrganization(newOrgName);
            selectOrganization(result.orgId);

            // Check if subscription setup is required
            if (result.requiresSubscription) {
                navigate(`/subscription-setup?orgId=${result.orgId}`, { replace: true });
            } else {
                navigate('/onboarding', { replace: true });
            }
        } catch (error) {
            console.error('Failed to create organization:', error);
            alert('Failed to create organization. Please try again.');
        } finally {
            setIsCreating(false);
        }
    };

    const handleJoinWithInvite = async () => {
        if (!inviteCode.trim() || isJoining) return;

        setIsJoining(true);
        try {
            const orgId = await joinOrganization(inviteCode);
            selectOrganization(orgId);
            navigate('/employee-questionnaire', { replace: true });
        } catch (error) {
            console.error('Failed to join organization:', error);
            alert('Failed to join organization. Please check your invite code.');
        } finally {
            setIsJoining(false);
        }
    };

    if (loading) {
        return <div className="p-8">Loading your organizations...</div>;
    }

    return (
        <div className="min-h-screen bg-[--background-primary] py-12">
            <div className="max-w-md mx-auto">
                <div className="text-center mb-8">
                    <div className="w-16 h-16 bg-[--accent] rounded-full flex items-center justify-center font-bold text-white text-2xl mx-auto mb-4">
                        A
                    </div>
                    <h1 className="text-3xl font-bold text-[--text-primary] mb-2">Welcome to Auditly</h1>
                    <p className="text-[--text-secondary]">Select an organization to continue</p>
                </div>

                {/* Existing Organizations */}
                {organizations.length > 0 && (
                    <Card className="mb-6">
                        <h2 className="text-lg font-semibold mb-4 text-[--text-primary]">Your Organizations</h2>
                        <div className="space-y-3">
                            {organizations.map((org) => (
                                <div
                                    key={org.orgId}
                                    className="flex items-center justify-between p-3 border border-[--border-color] rounded-lg hover:bg-[--background-secondary] cursor-pointer"
                                    onClick={() => handleSelectOrg(org.orgId)}
                                >
                                    <div>
                                        <div className="font-medium text-[--text-primary]">{org.name}</div>
                                        <div className="text-sm text-[--text-secondary]">
                                            {org.role} • {org.onboardingCompleted ? 'Active' : 'Setup Required'}
                                        </div>
                                    </div>
                                    <Button variant="secondary" size="sm">
                                        Enter
                                    </Button>
                                </div>
                            ))}
                        </div>
                    </Card>
                )}

                {/* Join with Invite */}
                {inviteCode && (
                    <Card className="mb-6">
                        <h2 className="text-lg font-semibold mb-4 text-[--text-primary]">Join Organization</h2>
                        <p className="text-sm text-[--text-secondary] mb-4">
                            You've been invited to join an organization with code: <code className="bg-[--background-tertiary] px-2 py-1 rounded text-[--text-primary]">{inviteCode}</code>
                        </p>
                        <Button onClick={handleJoinWithInvite} className="w-full" disabled={isJoining}>
                            {isJoining ? 'Joining...' : 'Accept Invitation'}
                        </Button>
                    </Card>
                )}

                {/* Create New Organization */}
                <Card>
                    <h2 className="text-lg font-semibold mb-4 text-[--text-primary]">Create New Organization</h2>
                    {!showCreateOrg ? (
                        <Button
                            onClick={() => setShowCreateOrg(true)}
                            variant="secondary"
                            className="w-full"
                        >
                            + Create New Organization
                        </Button>
                    ) : (
                        <div className="space-y-4">
                            <div>
                                <label className="block text-sm font-medium text-[--text-primary] mb-2">
                                    Organization Name
                                </label>
                                <input
                                    type="text"
                                    value={newOrgName}
                                    onChange={(e) => setNewOrgName(e.target.value)}
                                    placeholder="Enter organization name"
                                    className="w-full px-3 py-2 border border-[--input-border] rounded-lg bg-[--input-bg] text-[--text-primary] placeholder-[--input-placeholder] focus:outline-none focus:ring-2 focus:ring-[--accent] focus:border-[--accent]"
                                />
                            </div>
                            <div className="flex space-x-3">
                                <Button
                                    onClick={handleCreateOrg}
                                    disabled={!newOrgName.trim() || isCreating}
                                    className="flex-1"
                                >
                                    {isCreating ? 'Creating...' : 'Create Organization'}
                                </Button>
                                <Button
                                    onClick={() => setShowCreateOrg(false)}
                                    variant="secondary"
                                    className="flex-1"
                                >
                                    Cancel
                                </Button>
                            </div>
                        </div>
                    )}
                </Card>

                {/* Manual Invite Code Entry */}
                {!inviteCode && (
                    <Card className="mt-6">
                        <h2 className="text-lg font-semibold mb-4 text-[--text-primary]">Have an invite code?</h2>
                        <div className="flex space-x-3">
                            <input
                                type="text"
                                value={inviteCode}
                                onChange={(e) => setInviteCode(e.target.value)}
                                placeholder="Enter invite code"
                                className="flex-1 px-3 py-2 border border-[--input-border] rounded-lg bg-[--input-bg] text-[--text-primary] placeholder-[--input-placeholder] focus:outline-none focus:ring-2 focus:ring-[--accent] focus:border-[--accent]"
                            />
                            <Button
                                onClick={handleJoinWithInvite}
                                disabled={!inviteCode.trim() || isJoining}
                                variant="secondary"
                            >
                                {isJoining ? 'Joining...' : 'Join'}
                            </Button>
                        </div>
                    </Card>
                )}
            </div>
        </div>
    );
};

export default OrgSelection;