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

'use client';

import React, { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import { useTranslations } from 'next-intl';
import { useRouter } from '@/i18n/navigation';
import { authApi } from '@/lib/api/auth';
import {
  clearRememberedEmail,
  getRememberedEmail,
  setRememberedEmail,
} from '@/lib/auth/rememberMe';
import { User, Organization } from '@/types/organization';
import { isAppLocale, getLocaleFromPathname } from '@/i18n/routing';
import {
  clearAccessTokenExpiresAt,
  rememberAccessTokenExpiresAt,
  startProactiveSessionRefresh,
} from '@/lib/auth/proactiveRefresh';
import { asApiError, legacyStatusCode, type ApiError } from '@/types/api';

function toApiError(err: unknown): ApiError {
  return (
    asApiError(err) ?? {
      statusCode: legacyStatusCode(err) ?? 500,
      code:
        legacyStatusCode(err) === 401
          ? 'AUTH_UNAUTHORIZED'
          : legacyStatusCode(err) === 403
            ? 'PERMISSION_DENIED'
            : 'INTERNAL_ERROR',
    }
  );
}

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

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

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const t = useTranslations('auth');
  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 [apiError, setApiError] = useState<ApiError | 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),
          language: (payload.user as User).language ?? 'en',
        },
        organizations,
      };
    }

    if (payload?.id && payload?.email && payload?.name) {
      return {
        user: {
          id: payload.id,
          email: payload.email,
          name: payload.name,
          language: payload.language ?? 'en',
        },
        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);
        rememberAccessTokenExpiresAt(
          (response.data as { accessTokenExpiresAt?: string }).accessTokenExpiresAt,
        );

        setUser(userData);
        setOrganizations(orgs);

        const storedOrgId = localStorage.getItem('currentOrganizationId');
        if (storedOrgId && orgs.length > 0) {
          const org = orgs.find(o => o.id === storedOrgId);
          if (org) {
            const selected = await authApi.selectOrganization(org.id);
            setCurrentOrganization({
              id: selected.data.organization.id,
              name: selected.data.organization.name,
              type: selected.data.organization.type as Organization['type'],
              isOwner: Boolean(selected.data.organization.isOwner),
              permissions: selected.data.organization.permissions,
              plan: selected.data.organization.plan,
            });
          } else {
            setCurrentOrganization(null);
          }
        } else if (orgs.length === 1 && userData) {
          const selected = await authApi.selectOrganization(orgs[0].id);
          localStorage.setItem('currentOrganizationId', orgs[0].id);
          setCurrentOrganization({
            id: selected.data.organization.id,
            name: selected.data.organization.name,
            type: selected.data.organization.type as Organization['type'],
            isOwner: Boolean(selected.data.organization.isOwner),
            permissions: selected.data.organization.permissions,
            plan: selected.data.organization.plan,
          });
        } else {
          setCurrentOrganization(null);
        }
      }
    } catch (err: unknown) {
      const status = legacyStatusCode(err);

      // Access token may have expired while refresh cookie is still valid (e.g. JWT_EXPIRES_IN=15m).
      if (status === 401) {
        try {
          const refreshResult = await authApi.refreshSessionFromCookies();
          rememberAccessTokenExpiresAt(refreshResult.data?.accessTokenExpiresAt);
          const orgId = localStorage.getItem('currentOrganizationId');
          if (orgId) {
            await authApi.selectOrganization(orgId);
          }
          const retry = await authApi.getProfile();
          if (retry.success) {
            const { user: userData, organizations: orgs } = normalizeProfilePayload(retry.data);
            rememberAccessTokenExpiresAt(
              (retry.data as { accessTokenExpiresAt?: string }).accessTokenExpiresAt,
            );
            setUser(userData);
            setOrganizations(orgs);

            const storedOrgId = localStorage.getItem('currentOrganizationId');
            if (storedOrgId && orgs.length > 0) {
              const org = orgs.find((o) => o.id === storedOrgId);
              if (org) {
                const selected = await authApi.selectOrganization(org.id);
                setCurrentOrganization({
                  id: selected.data.organization.id,
                  name: selected.data.organization.name,
                  type: selected.data.organization.type as Organization['type'],
                  isOwner: Boolean(selected.data.organization.isOwner),
                  permissions: selected.data.organization.permissions,
                  plan: selected.data.organization.plan,
                });
              } else {
                setCurrentOrganization(null);
              }
            } else if (orgs.length === 1 && userData) {
              const selected = await authApi.selectOrganization(orgs[0].id);
              localStorage.setItem('currentOrganizationId', orgs[0].id);
              setCurrentOrganization({
                id: selected.data.organization.id,
                name: selected.data.organization.name,
                type: selected.data.organization.type as Organization['type'],
                isOwner: Boolean(selected.data.organization.isOwner),
                permissions: selected.data.organization.permissions,
                plan: selected.data.organization.plan,
              });
            } else {
              setCurrentOrganization(null);
            }
            return;
          }
        } catch {
          /* fall through to logged-out state */
        }
      }

      // 401 on profile is expected when there is no session — not an application error.
      if (status !== 401) {
        console.error('Auth check failed:', err);
      }

      // Only clear state — DO NOT redirect here
      setUser(null);
      setOrganizations([]);
      setCurrentOrganization(null);
      clearAccessTokenExpiresAt();
    } finally {
      setIsLoading(false);
      setIsAuthReady(true);
    }
  }, [normalizeProfilePayload]);

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

  useEffect(() => {
    if (!user || !isAuthReady) {
      return;
    }

    return startProactiveSessionRefresh();
  }, [user, isAuthReady]);

  const applyUrlLocaleToUser = useCallback(async (user: User): Promise<User> => {
    if (typeof window === 'undefined') return user;

    const urlLocale = getLocaleFromPathname(window.location.pathname);
    if (!isAppLocale(urlLocale) || user.language === urlLocale) {
      return user;
    }

    try {
      await authApi.updateLanguage(urlLocale);
    } catch {
      /* keep URL locale in client state even if persistence fails */
    }

    return { ...user, language: urlLocale };
  }, []);

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

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

      const userData = await applyUrlLocaleToUser(response.data.user);
      rememberAccessTokenExpiresAt(response.data.accessTokenExpiresAt);
      setUser(userData);
      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) {
      setApiError(toApiError(err));
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [applyUrlLocaleToUser, router, t]);

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

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

      if (rememberMe) {
        setRememberedEmail(email);
      } else {
        clearRememberedEmail();
      }

      const userData = await applyUrlLocaleToUser(response.data.user);
      rememberAccessTokenExpiresAt(response.data.accessTokenExpiresAt);
      setUser(userData);
      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) {
      setApiError(toApiError(err));
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [applyUrlLocaleToUser, router, t]);

  const logout = useCallback(async () => {
    try {
      // Important: clear auth cookies/session on the server first,
      // otherwise proxy may still treat the user as authenticated.
      await authApi.logout();
    } catch (err) {
      console.error('Logout API failed:', err);
    } finally {
      const rememberedEmail = getRememberedEmail();
      localStorage.clear();
      if (rememberedEmail) {
        setRememberedEmail(rememberedEmail);
      }
      setUser(null);
      setOrganizations([]);
      setCurrentOrganization(null);
      setApiError(null);
      clearAccessTokenExpiresAt();
      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,
      });

      const redirectPath =
        typeof window !== 'undefined'
          ? sessionStorage.getItem('authRedirect')
          : null;
      if (redirectPath) {
        sessionStorage.removeItem('authRedirect');
        router.push(redirectPath);
      } else {
        router.push('/today');
      }

    } catch (err: any) {
      setApiError(toApiError(err));
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [router]);

  const createOrganization = useCallback(async (
    organizationName: string,
    organizationEmail: string,
    organizationType: 'CLINIC' | 'LAB',
    planName?: string,
  ) => {
    try {
      setIsLoading(true);
      setApiError(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) {
      setApiError(toApiError(err));
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [normalizeProfilePayload, t]);

  const refreshSession = useCallback(async () => {
    await checkAuth();
  }, [checkAuth]);

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

  const setUserLanguage = useCallback((language: string) => {
    const normalized = isAppLocale(language) ? language : 'en';
    setUser((current) => (current ? { ...current, language: normalized } : current));
  }, []);

  const contextValue = useMemo(
    () => ({
      user,
      organizations,
      currentOrganization,
      isLoading,
      isAuthReady,
      apiError,
      registerTrial,
      login,
      logout,
      selectOrganization,
      createOrganization,
      setUserLanguage,
      refreshSession,
      clearError,
    }),
    [
      user,
      organizations,
      currentOrganization,
      isLoading,
      isAuthReady,
      apiError,
      registerTrial,
      login,
      logout,
      selectOrganization,
      createOrganization,
      setUserLanguage,
      refreshSession,
      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;
};