popiart-node-canvas/src/hooks/useSelectedModelSchema.ts

"use client";

import { useEffect, useRef, useState } from "react";
import { ModelInputDef, SelectedModel } from "@/types";
import { ModelParameter, ProviderModel } from "@/lib/providers/types";
import { useProviderApiKeys } from "@/store/workflowStore";

interface ModelSchemaResult {
  parameters: ModelParameter[];
  inputs: ModelInputDef[];
  model?: ProviderModel;
}

interface UseSelectedModelSchemaOptions {
  selectedModel?: SelectedModel;
  enabled?: boolean;
  refreshKey?: number;
  onLoaded?: (schema: ModelSchemaResult) => void;
}

export function useSelectedModelSchema({
  selectedModel,
  enabled = true,
  refreshKey,
  onLoaded,
}: UseSelectedModelSchemaOptions) {
  const [parameters, setParameters] = useState<ModelParameter[]>([]);
  const [inputs, setInputs] = useState<ModelInputDef[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const requestKeyRef = useRef<string | null>(null);
  const abortRef = useRef<AbortController | null>(null);
  const onLoadedRef = useRef(onLoaded);
  const {
    newApiWGBaseUrl,
    geminiApiKey,
    replicateApiKey,
    falApiKey,
    kieApiKey,
    wavespeedApiKey,
  } = useProviderApiKeys();

  useEffect(() => {
    onLoadedRef.current = onLoaded;
  }, [onLoaded]);

  useEffect(() => {
    const provider = selectedModel?.provider;
    const modelId = selectedModel?.modelId;
    if (!enabled || !provider || !modelId) {
      if (requestKeyRef.current) {
        requestKeyRef.current = null;
        abortRef.current?.abort();
        setParameters([]);
        setInputs([]);
        setError(null);
        setIsLoading(false);
        onLoadedRef.current?.({ parameters: [], inputs: [] });
      }
      return;
    }

    const requestKey = `${provider}:${modelId}:${refreshKey ?? 0}`;
    if (requestKeyRef.current === requestKey) return;
    requestKeyRef.current = requestKey;
    abortRef.current?.abort();

    const controller = new AbortController();
    abortRef.current = controller;
    setIsLoading(true);
    setError(null);

    const headers: HeadersInit = {};
    if (newApiWGBaseUrl) headers["X-NewApiWG-Base-URL"] = newApiWGBaseUrl;
    if (geminiApiKey) headers["X-Gemini-API-Key"] = geminiApiKey;
    if (replicateApiKey) headers["X-Replicate-Key"] = replicateApiKey;
    if (falApiKey) headers["X-Fal-Key"] = falApiKey;
    if (kieApiKey) headers["X-Kie-Key"] = kieApiKey;
    if (wavespeedApiKey) headers["X-WaveSpeed-Key"] = wavespeedApiKey;

    fetch(`/api/models/${encodeURIComponent(modelId)}?provider=${provider}`, {
      headers,
      signal: controller.signal,
    })
      .then(async (response) => {
        const data = await response.json().catch(() => ({}));
        if (!response.ok) {
          throw new Error(data.error || `Failed to fetch schema: ${response.status}`);
        }
        if (controller.signal.aborted || requestKeyRef.current !== requestKey) {
          return;
        }
        const nextSchema = {
          parameters: data.parameters || [],
          inputs: data.inputs || [],
          model: data.model,
        };
        setParameters(nextSchema.parameters);
        setInputs(nextSchema.inputs);
        onLoadedRef.current?.(nextSchema);
      })
      .catch((err) => {
        if (controller.signal.aborted || requestKeyRef.current !== requestKey) return;
        setParameters([]);
        setInputs([]);
        setError(err instanceof Error ? err.message : "Failed to fetch schema");
      })
      .finally(() => {
        if (!controller.signal.aborted && requestKeyRef.current === requestKey) {
          setIsLoading(false);
        }
      });

    return () => {
      if (abortRef.current === controller) {
        controller.abort();
        abortRef.current = null;
        requestKeyRef.current = null;
      } else {
        controller.abort();
      }
    };
  }, [
    enabled,
    selectedModel?.provider,
    selectedModel?.modelId,
    refreshKey,
    newApiWGBaseUrl,
    geminiApiKey,
    replicateApiKey,
    falApiKey,
    kieApiKey,
    wavespeedApiKey,
  ]);

  return { parameters, inputs, isLoading, error };
}