dyolink/frontend/src/lib/hooks/useAuth.tsx

'use client';

import React, { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import { useRouter } from 'next/navigation';
import { authApi } from '@/lib/api/auth';
import { User, Organization } from '@/types/organization';

interface AuthContextType {
  user: User | null;
  organizations: Organization[];
  currentOrganization: Organization | null;
  isLoading: boolean;
  isAuthReady: boolean; // ✅ NEW
  error: string | null;
  registerTrial: (
    email: string,
    password: string,
    name: string,
    organizationName: string,
    organizationEmail: string,
    organizationType: 'CLINIC' | 'LAB'
  ) => Promise<void>;
  login: (email: string, password: string) => Promise<void>;
  logout: () => Promise<void>;
  selectOrganization: (orgId: string) => Promise<void>;
  createOrganization: (
    organizationName: string,
    organizationEmail: string,
    organizationType: 'CLINIC' | 'LAB',
    planName?: string,
  ) => Promise<string>;
  clearError: () => void;
}

const AuthContext = createContext<AuthContextType | undefined>(undefined);

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const [user, setUser] = useState<User | null>(null);
  const [organizations, setOrganizations] = useState<Organization[]>([]);
  const [currentOrganization, setCurrentOrganization] = useState<Organization | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [isAuthReady, setIsAuthReady] = useState(false); // ✅ KEY FIX
  const [error, setError] = useState<string | null>(null);

  const router = useRouter();

  const normalizeProfilePayload = useCallback((payload: any): { user: User | null; organizations: Organization[] } => {
    const organizations = payload?.organizations || [];

    if (payload?.user) {
      return { user: payload.user as User, organizations };
    }

    if (payload?.id && payload?.email && payload?.name) {
      return {
        user: {
          id: payload.id,
          email: payload.email,
          name: payload.name,
        },
        organizations,
      };
    }

    return { user: null, organizations };
  }, []);

  const checkAuth = useCallback(async () => {
    try {
      setIsLoading(true);

      const response = await authApi.getProfile();

      if (response.success) {
        const { user: userData, organizations: orgs } = normalizeProfilePayload(response.data);

        setUser(userData);
        setOrganizations(orgs);

        const storedOrgId = localStorage.getItem('currentOrganizationId');
        if (storedOrgId && orgs.length > 0) {
          const org = orgs.find(o => o.id === storedOrgId);
          if (org) {
            setCurrentOrganization(org);
            // Ensure cookie token carries organizationId for org-scoped APIs.
            await authApi.selectOrganization(org.id);
          } else {
            setCurrentOrganization(null);
          }
        } else if (orgs.length === 1 && userData) {
          setCurrentOrganization(orgs[0]);
          localStorage.setItem('currentOrganizationId', orgs[0].id);
          // Keep JWT in sync with selected org even for single-org users.
          await authApi.selectOrganization(orgs[0].id);
        } else {
          setCurrentOrganization(null);
        }
      }
    } catch (err) {
      console.error('Auth check failed:', err);
      // Only clear state — DO NOT redirect here
      setUser(null);
      setOrganizations([]);
      setCurrentOrganization(null);
    } finally {
      setIsLoading(false);
      setIsAuthReady(true);
    }
  }, [normalizeProfilePayload]);

  useEffect(() => {
    void checkAuth();
  }, [checkAuth]);

  // ✅ REGISTER
  const registerTrial = useCallback(async (
    email: string,
    password: string,
    name: string,
    organizationName: string,
    organizationEmail: string,
    organizationType: 'CLINIC' | 'LAB'
  ) => {
    try {
      setIsLoading(true);
      setError(null);

      const response = await authApi.registerTrial({
        email,
        password,
        name,
        organizationName,
        organizationEmail,
        organizationType,
      });

      setUser(response.data.user);
      setOrganizations(response.data.organizations);

      const orgs = response.data.organizations;

      if (orgs.length === 1) {
        const org = orgs[0];
        await authApi.selectOrganization(org.id);
        setCurrentOrganization(org);
        localStorage.setItem('currentOrganizationId', org.id);
        router.push('/today');
      } else {
        router.push('/select-organization');
        
      }

    } catch (err: any) {
      setError(err.message || 'Registration failed');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [router]);

  // ✅ LOGIN
  const login = useCallback(async (email: string, password: string) => {
    try {
      setIsLoading(true);
      setError(null);

      const response = await authApi.login({ email, password });

      setUser(response.data.user);
      setOrganizations(response.data.organizations);

      const orgs = response.data.organizations;

      if (orgs.length === 1) {
        const org = orgs[0];
        await authApi.selectOrganization(org.id);
        setCurrentOrganization(org);
        localStorage.setItem('currentOrganizationId', org.id);
        router.push('/today');
      } else {
        router.push('/select-organization');
      }

    } catch (err: any) {
      setError(err.message || 'Login failed');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [router]);

  const logout = useCallback(async () => {
    try {
      // Important: clear auth cookies/session on the server first,
      // otherwise middleware may still treat the user as authenticated.
      await authApi.logout();
    } catch (err) {
      console.error('Logout API failed:', err);
    } finally {
      localStorage.clear();
      setUser(null);
      setOrganizations([]);
      setCurrentOrganization(null);
      setError(null);
      setIsAuthReady(true);
      router.replace('/');
      router.refresh();
    }
  }, [router]);

  const selectOrganization = useCallback(async (orgId: string) => {
    try {
      setIsLoading(true);

      const response = await authApi.selectOrganization(orgId);

      const { organization } = response.data;

      localStorage.setItem('currentOrganizationId', organization.id);

      setCurrentOrganization({
        id: organization.id,
        name: organization.name,
        type: organization.type as Organization['type'],
        isOwner: Boolean((organization as { isOwner?: boolean }).isOwner),
        permissions: (organization as { permissions?: string[] }).permissions,
        plan: (organization as { plan?: Organization['plan'] }).plan,
      });

      router.push('/today');

    } catch (err: any) {
      setError(err.message);
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [router]);

  const createOrganization = useCallback(async (
    organizationName: string,
    organizationEmail: string,
    organizationType: 'CLINIC' | 'LAB',
    planName?: string,
  ) => {
    try {
      setIsLoading(true);
      setError(null);

      const createResponse = await authApi.createOrganization({
        organizationName,
        organizationEmail,
        organizationType,
        planName,
      });

      const profileResponse = await authApi.getProfile();
      if (profileResponse.success) {
        const { user: userData, organizations: orgs } = normalizeProfilePayload(profileResponse.data);
        setUser(userData);
        setOrganizations(orgs);
      }

      return createResponse.data.organization.id as string;
    } catch (err: any) {
      setError(err.message || 'Failed to create organization');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [normalizeProfilePayload]);

  const clearError = useCallback(() => setError(null), []);

  const contextValue = useMemo(
    () => ({
      user,
      organizations,
      currentOrganization,
      isLoading,
      isAuthReady,
      error,
      registerTrial,
      login,
      logout,
      selectOrganization,
      createOrganization,
      clearError,
    }),
    [
      user,
      organizations,
      currentOrganization,
      isLoading,
      isAuthReady,
      error,
      registerTrial,
      login,
      logout,
      selectOrganization,
      createOrganization,
      clearError,
    ],
  );

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

export const useAuth = () => {
  const context = useContext(AuthContext);
  if (!context) throw new Error('useAuth must be used within AuthProvider');
  return context;
};