UHD_Media_Tool/apps/web/app/api/pi.ts

import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import type {
  PiModelsProbeRequest,
  PiOAuthCompleteRequest,
  PiOAuthCompleteResult,
  PiOAuthStartRequest,
  PiOAuthStartResult,
  PiProvidersResponse,
  PiUpstream,
  PiUpstreamCreate,
  PiUpstreamList,
  PiUpstreamModels,
  PiUpstreamTestRequest,
  PiUpstreamTestResult,
  PiUpstreamUpdate,
} from "@umt/contracts";
import { apiData, TTL_MINUTE } from "~/lib/api.ts";
import { qk } from "~/lib/query-keys.ts";

// Agent 上游(W7-A 的 /pi/*):能力表、上游 CRUD、测试、模型表、订阅登录。
// 写操作后刷新上游列表与全局设置(默认上游存在 settings.pi.defaultUpstreamId)。

export const piKeys = {
  all: ["pi"] as const,
  providers: (withModels: boolean) => ["pi", "providers", withModels] as const,
  upstreams: ["pi", "upstreams"] as const,
  models: (id: string) => ["pi", "models", id] as const,
};

/** 能力表;withModels = 带静态模型表(约 370 KB,只在新建向导里要) */
export function usePiProviders(withModels = false, enabled = true) {
  return useQuery({
    queryKey: piKeys.providers(withModels),
    enabled,
    staleTime: 30 * TTL_MINUTE,
    queryFn: ({ signal }) =>
      apiData<PiProvidersResponse>("/pi/providers", { params: withModels ? {} : { models: "false" }, signal }),
  });
}

export function usePiUpstreams() {
  return useQuery({
    queryKey: piKeys.upstreams,
    queryFn: ({ signal }) => apiData<PiUpstreamList>("/pi/upstreams", { signal }),
  });
}

export function usePiUpstreamModels(id: string | null) {
  return useQuery({
    queryKey: piKeys.models(id ?? ""),
    enabled: !!id,
    staleTime: 5 * TTL_MINUTE,
    queryFn: ({ signal }) =>
      apiData<PiUpstreamModels>("/pi/upstreams/{id}/models", { pathParams: { id: id ?? "" }, signal }),
  });
}

/** 获取模型列表:用该上游的凭证调 provider 的模型列表接口;失败时响应里仍是当前列表(静态表)并带 error */
export function useRefreshPiModels() {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (id: string) =>
      apiData<PiUpstreamModels>("/pi/upstreams/{id}/models/refresh", { method: "POST", pathParams: { id }, body: {} }),
    onSuccess: (data, id) => qc.setQueryData(piKeys.models(id), data),
  });
}

/** 新建上游前试拉模型列表(凭证只用这一次,不落库) */
export function useProbePiModels() {
  return useMutation({
    mutationFn: (body: PiModelsProbeRequest) => apiData<PiUpstreamModels>("/pi/models/probe", { method: "POST", body }),
  });
}

function useInvalidatePi() {
  const qc = useQueryClient();
  return () => {
    qc.invalidateQueries({ queryKey: piKeys.upstreams });
    qc.invalidateQueries({ queryKey: qk.settings });
  };
}

export function useCreatePiUpstream() {
  const inv = useInvalidatePi();
  return useMutation({
    mutationFn: (body: PiUpstreamCreate) => apiData<PiUpstream>("/pi/upstreams", { method: "POST", body }),
    onSuccess: inv,
  });
}

export function useUpdatePiUpstream() {
  const inv = useInvalidatePi();
  return useMutation({
    mutationFn: ({ id, body }: { id: string; body: PiUpstreamUpdate }) =>
      apiData<PiUpstream>("/pi/upstreams/{id}", { method: "PATCH", pathParams: { id }, body }),
    onSuccess: inv,
  });
}

export function useRemovePiUpstream() {
  const inv = useInvalidatePi();
  return useMutation({
    mutationFn: (id: string) => apiData<unknown>("/pi/upstreams/{id}", { method: "DELETE", pathParams: { id } }),
    onSuccess: inv,
  });
}

export function useTestPiUpstream() {
  const inv = useInvalidatePi();
  return useMutation({
    mutationFn: ({ id, body }: { id: string; body: PiUpstreamTestRequest }) =>
      apiData<PiUpstreamTestResult>("/pi/upstreams/{id}/test", { method: "POST", pathParams: { id }, body }),
    onSettled: inv,
  });
}

export function useStartPiOAuth() {
  return useMutation({
    mutationFn: ({ id, body }: { id: string; body: PiOAuthStartRequest }) =>
      apiData<PiOAuthStartResult>("/pi/upstreams/{id}/oauth/start", { method: "POST", pathParams: { id }, body }),
  });
}

/** 完成 / 轮询一次登录(device 模式前端反复调用直到终态) */
export function completePiOAuth(id: string, body: PiOAuthCompleteRequest, signal?: AbortSignal) {
  return apiData<PiOAuthCompleteResult>("/pi/upstreams/{id}/oauth/complete", {
    method: "POST",
    pathParams: { id },
    body,
    signal,
  });
}

export function useLogoutPiOAuth() {
  const inv = useInvalidatePi();
  return useMutation({
    mutationFn: (id: string) =>
      apiData<PiUpstream>("/pi/upstreams/{id}/oauth/logout", { method: "POST", pathParams: { id }, body: {} }),
    onSuccess: inv,
  });
}

export { useInvalidatePi };