"use client";

/**
 * OAuth2 PKCE flow for the *user* connect path.
 *
 * Unlike useMcpOAuthFlow (used in the admin create-server form), this hook
 * targets a server that already exists in the database.  It therefore skips
 * the temp-session cache step and calls /server/oauth/{serverId}/authorize
 * directly with the real server_id.
 *
 * On success it calls storeMCPOAuthUserCredential to persist the token for
 * the user and then invokes onSuccess so the caller can refresh UI state.
 */

import { useCallback, useEffect, useRef, useState } from "react";
import {
  buildMcpOAuthAuthorizeUrl,
  exchangeMcpOAuthToken,
  registerMcpOAuthClient,
  storeMCPOAuthUserCredential,
} from "@/components/networking";
import { toast } from "@/lib/toast";
import { extractErrorMessage } from "@/utils/errorUtils";
import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils";

export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";

interface UseUserMcpOAuthFlowOptions {
  accessToken: string;
  serverId: string;
  serverAlias?: string | null;
  /** Scopes to request, e.g. ["repo", "read:user"] */
  scopes?: string[];
  /** Pre-configured client_id if the MCP server record has one. */
  clientId?: string | null;
  onSuccess: () => void;
}

interface UseUserMcpOAuthFlowResult {
  startOAuthFlow: () => Promise<void>;
  status: UserMcpOAuthStatus;
  error: string | null;
}

const FLOW_STATE_KEY = "litellm-user-mcp-oauth-flow-state";
// Use a user-flow-specific key to avoid collisions with the admin OAuth flow
// (useMcpOAuthFlow) which uses "litellm-mcp-oauth-result".
const RESULT_KEY = "litellm-user-mcp-oauth-result";
const RETURN_URL_KEY = "litellm-mcp-oauth-return-url";

type StoredFlowState = {
  state: string;
  codeVerifier: string;
  serverId: string;
  redirectUri: string;
  clientId?: string;
  clientSecret?: string;
  scopes?: string[];
};

const setStorage = (key: string, value: string) => {
  setSecureItem(key, value);
};

const getStorage = (key: string): string | null => {
  return getSecureItem(key);
};

export const useUserMcpOAuthFlow = ({
  accessToken,
  serverId,
  serverAlias,
  scopes,
  clientId: preClientId,
  onSuccess,
}: UseUserMcpOAuthFlowOptions): UseUserMcpOAuthFlowResult => {
  const [status, setStatus] = useState<UserMcpOAuthStatus>("idle");
  const [error, setError] = useState<string | null>(null);
  const processingRef = useRef(false);

  const startOAuthFlow = useCallback(async () => {
    if (typeof window === "undefined") return;
    try {
      setStatus("authorizing");
      setError(null);

      let clientId: string | undefined = preClientId ?? undefined;
      let clientSecret: string | undefined;

      if (!clientId) {
        // Attempt dynamic client registration against the server's registration endpoint.
        try {
          const reg = await registerMcpOAuthClient(accessToken, serverId, {
            client_name: serverAlias || serverId,
            grant_types: ["authorization_code", "refresh_token"],
            response_types: ["code"],
            token_endpoint_auth_method: "none",
          });
          clientId = reg?.client_id;
          clientSecret = reg?.client_secret;
        } catch (_) {
          // Registration is optional; proceed without client_id
        }
      }

      const verifier = generateCodeVerifier();
      const challenge = await generateCodeChallenge(verifier);
      const state = crypto.randomUUID();
      const redirectUri = buildCallbackUrl();
      const scopeString = scopes?.filter((s) => s.trim()).join(" ");

      const authorizeUrl = buildMcpOAuthAuthorizeUrl({
        serverId,
        clientId,
        redirectUri,
        state,
        codeChallenge: challenge,
        scope: scopeString,
      });

      const flowState: StoredFlowState = {
        state,
        codeVerifier: verifier,
        serverId,
        redirectUri,
        clientId,
        clientSecret,
        scopes,
      };

      setStorage(FLOW_STATE_KEY, JSON.stringify(flowState));
      const returnUrl = new URL(window.location.href);
      returnUrl.searchParams.set("mcpOauthReturn", "apps");
      setStorage(RETURN_URL_KEY, returnUrl.toString());

      window.location.href = authorizeUrl;
    } catch (err) {
      const msg = extractErrorMessage(err);
      setError(msg);
      setStatus("error");
      toast.error(msg);
    }
  }, [accessToken, serverId, serverAlias, scopes, preClientId]);

  const resumeOAuthFlow = useCallback(async () => {
    if (typeof window === "undefined" || processingRef.current) return;

    const storedResult = getStorage(RESULT_KEY);
    if (!storedResult) return;

    // When multiple OAuth2ConnectButton components are mounted (one per server
    // card), each holds its own hook instance.  All run resumeOAuthFlow() on
    // mount and would compete for the same RESULT_KEY.  Peek at the stored
    // flow state first: only the hook instance whose serverId matches the one
    // that initiated the OAuth flow should consume the result.
    // Guard: only proceed if this hook's flow state exists (startOAuthFlow was
    // called from this hook).  Without the guard, a tools re-auth redirect writes
    // to the user result key too, and every OAuth2ConnectButton instance would try
    // to resume a flow that was never started here.
    const rawFlowState = getStorage(FLOW_STATE_KEY);
    if (!rawFlowState) return;

    try {
      const peeked = JSON.parse(rawFlowState) as StoredFlowState;
      if (peeked.serverId && peeked.serverId !== serverId) return;
    } catch (_) {}

    processingRef.current = true;
    clearStorage(RESULT_KEY);

    let payload: Record<string, unknown> | null = null;
    let flowState: StoredFlowState | null = null;

    try {
      payload = JSON.parse(storedResult);
      const raw = getStorage(FLOW_STATE_KEY);
      flowState = raw ? JSON.parse(raw) : null;
    } catch (_) {
      setError("Failed to resume OAuth flow. Please retry.");
      setStatus("error");
      processingRef.current = false;
      clearStorage(FLOW_STATE_KEY);
      return;
    }

    try {
      if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) {
        throw new Error("OAuth session state was lost. Please retry.");
      }
      if (!payload?.state || payload.state !== flowState.state) {
        throw new Error("OAuth state mismatch. Please retry.");
      }
      if (payload.error) {
        throw new Error((payload.error_description as string) || (payload.error as string));
      }
      if (!payload.code) {
        throw new Error("Authorization code missing in callback.");
      }

      setStatus("exchanging");
      const token = await exchangeMcpOAuthToken({
        serverId: flowState.serverId,
        code: payload.code as string,
        clientId: flowState.clientId,
        clientSecret: flowState.clientSecret,
        codeVerifier: flowState.codeVerifier,
        redirectUri: flowState.redirectUri,
        accessToken,
      });

      // Persist the token for this user via the backend.
      // accessToken comes from props — it is never stored in sessionStorage.
      await storeMCPOAuthUserCredential(accessToken, flowState.serverId, {
        access_token: token.access_token,
        refresh_token: token.refresh_token,
        expires_in: token.expires_in,
        scopes: flowState.scopes,
      });

      setStatus("success");
      setError(null);
      toast.success("Connected successfully");
      onSuccess();
    } catch (err) {
      const msg = extractErrorMessage(err);
      setError(msg);
      setStatus("error");
      toast.error(msg);
    } finally {
      clearStorage(FLOW_STATE_KEY);
      setTimeout(() => {
        processingRef.current = false;
      }, 1000);
    }
  }, [accessToken, serverId, onSuccess]);

  useEffect(() => {
    resumeOAuthFlow();
  }, [resumeOAuthFlow]);

  return { startOAuthFlow, status, error };
};
