import { createContext, useCallback, useContext, useMemo, useState } from 'react';
import type { ReactNode } from 'react';

const DEFAULT_AUTH_TOKEN_KEY = 'authToken';
const DEFAULT_REFRESH_TOKEN_KEY = 'refreshToken';
const DEFAULT_LEGACY_KEYS = ['auth_token', 'auth_user', 'token'];

export type AuthState<TUser> = {
    authToken: string | null;
    refreshToken: string | null;
    currentUser: TUser | null;
};

export type AuthContextValue<TUser> = AuthState<TUser> & {
    setSession: (authToken: string, refreshToken: string, currentUser: TUser) => void;
    setCurrentUser: (currentUser: TUser | null) => void;
    clearSession: () => void;
};

export type CreateAuthContextOptions = {
    authTokenKey?: string;
    refreshTokenKey?: string;
    legacyKeys?: string[];
};

export function createAuthContext<TUser>(options: CreateAuthContextOptions = {}) {
    const authTokenKey = options.authTokenKey ?? DEFAULT_AUTH_TOKEN_KEY;
    const refreshTokenKey = options.refreshTokenKey ?? DEFAULT_REFRESH_TOKEN_KEY;
    const legacyKeys = options.legacyKeys ?? DEFAULT_LEGACY_KEYS;

    const AuthContext = createContext<AuthContextValue<TUser> | undefined>(undefined);

    function readStoredSession(): AuthState<TUser> {
        for (const key of legacyKeys) {
            localStorage.removeItem(key);
        }

        const authToken = localStorage.getItem(authTokenKey);
        const refreshToken = localStorage.getItem(refreshTokenKey);

        return {
            authToken,
            refreshToken,
            currentUser: null,
        };
    }

    function AuthProvider({ children }: Readonly<{ children: ReactNode }>) {
        const [state, setState] = useState<AuthState<TUser>>(readStoredSession);

        const setSession = useCallback(
            (authToken: string, refreshToken: string, currentUser: TUser) => {
                localStorage.setItem(authTokenKey, authToken);
                localStorage.setItem(refreshTokenKey, refreshToken);
                setState({ authToken, refreshToken, currentUser });
            },
            [],
        );

        const clearSession = useCallback(() => {
            localStorage.removeItem(authTokenKey);
            localStorage.removeItem(refreshTokenKey);
            setState({ authToken: null, refreshToken: null, currentUser: null });
        }, []);

        const setCurrentUser = useCallback((currentUser: TUser | null) => {
            setState((prev) => ({ ...prev, currentUser }));
        }, []);

        const value = useMemo<AuthContextValue<TUser>>(
            () => ({
                ...state,
                setSession,
                setCurrentUser,
                clearSession,
            }),
            [state, setSession, setCurrentUser, clearSession],
        );

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

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

    return {
        AuthProvider,
        useAuth,
        AuthContext,
    };
}