frontend/src/components/HuggingFaceLoginButton.tsx

import { Tooltip } from "@mui/material";
import React, { useCallback, useEffect } from "react";
import { UserInfo } from "../types";
import {
  exchangeCodeForToken,
  fetchUserInfo,
  readFragmentParams,
  startLogin,
} from "../utils/oauth";
import { LoginButton } from "./LoginButton";

interface HuggingFaceLoginButtonProps {
  userInfo: UserInfo | null;
  accessToken: string | null;
  loginLabel: string;
  isDisabled?: boolean;
  onLoginStateChange: (
    userInfo: UserInfo | null,
    accessToken: string | null,
    loginLabel: string,
  ) => void;
}

export const HuggingFaceLoginButton: React.FC<HuggingFaceLoginButtonProps> = ({
  userInfo,
  accessToken,
  loginLabel,
  isDisabled = false,
  onLoginStateChange,
}) => {
  const isLoggedIn = !!userInfo?.sub;
  // Handle OAuth redirect
  const handleRedirect = useCallback(async () => {
    const params = new URLSearchParams(window.location.search);
    const { access_token: fragToken, error: fragErr } = readFragmentParams();

    if (fragErr) {
      onLoginStateChange(null, null, `Error: ${fragErr}`);
      return true;
    }

    const error = params.get("error");
    const errorDescription = params.get("error_description");
    if (error) {
      onLoginStateChange(
        null,
        null,
        `Error: ${error}${errorDescription ? ` — ${errorDescription}` : ""}`,
      );
      return true;
    }

    const returnedState = params.get("state");
    const expectedState = sessionStorage.getItem("hf_oauth_state");
    if (returnedState && expectedState && returnedState !== expectedState) {
      onLoginStateChange(null, null, "Error: invalid state");
      return true;
    }

    // Implicit flow
    if (fragToken) {
      try {
        const info = await fetchUserInfo(fragToken);
        const label =
          info?.email || info?.name || info?.preferred_username || "User";
        onLoginStateChange(info, fragToken, label);
      } catch (err) {
        console.error(err);
        onLoginStateChange(null, fragToken, "Connected");
      }
      window.history.replaceState({}, "", window.location.pathname);
      return true;
    }

    const code = params.get("code");
    if (!code) return false;

    try {
      const tokenResponse = await exchangeCodeForToken(code);
      const token = tokenResponse.access_token;
      if (token) {
        try {
          const info = await fetchUserInfo(token);
          const label =
            info?.email || info?.name || info?.preferred_username || "User";
          onLoginStateChange(info, token, label);
        } catch (err) {
          console.error(err);
          onLoginStateChange(null, token, "Connected");
        }
      } else {
        onLoginStateChange(null, null, "Connected");
      }
    } catch (e: any) {
      console.error(e);
      onLoginStateChange(null, null, `Authentication error: ${e.message}`);
    } finally {
      window.history.replaceState({}, "", window.location.pathname);
    }

    return true;
  }, [onLoginStateChange]);

  // Initialize on component mount
  useEffect(() => {
    const initialize = async () => {
      const handled = await handleRedirect();
      if (!handled) {
        onLoginStateChange(null, null, "Login with Hugging Face");
      }
    };
    initialize();
  }, [handleRedirect, onLoginStateChange]);

  const handleLoginClick = () => {
    onLoginStateChange(userInfo, accessToken, "Redirecting to Hugging Face…");
    startLogin().catch((err) => {
      console.error(err);
      onLoginStateChange(userInfo, accessToken, `Error: ${err.message}`);
    });
  };

  return (
    <Tooltip title={isLoggedIn ? "Log out" : ""} placement="right">
      <LoginButton
        icon={
          <img
            src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
            alt="Hugging Face"
            style={{ width: 18, height: 18 }}
          />
        }
        onClick={
          isLoggedIn
            ? () => onLoginStateChange(null, null, "Login with Hugging Face")
            : handleLoginClick
        }
        isLoggedIn={isLoggedIn}
        disabled={isDisabled}
      >
        {loginLabel}
      </LoginButton>
    </Tooltip>
  );
};