ai-chatbot/components/model-selector.tsx

"use client";

import type { Session } from "next-auth";
import {
  startTransition,
  useEffect,
  useMemo,
  useOptimistic,
  useState,
} from "react";
import useSWR from "swr";
import { saveChatModelAsCookie } from "@/app/(chat)/actions";
import { Button } from "@/components/ui/button";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuSeparator,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { cn } from "@/lib/utils";
import { CheckCircleFillIcon, ChevronDownIcon } from "@/components/icons";

// 从 API 获取的动态模型类型
type DynamicModel = {
  id: string;
  modelId: string;
  displayName: string;
  description: string | null;
  supportsVision: boolean;
  supportsImageGen: boolean;
  isDefault: boolean;
  providerName: string;
  providerType: string;
};

const fetcher = (url: string) => fetch(url).then((res) => res.json());

export function ModelSelector({
  session,
  selectedModelId,
  className,
}: {
  session: Session;
  selectedModelId: string;
} & React.ComponentProps<typeof Button>) {
  const [open, setOpen] = useState(false);
  const [optimisticModelId, setOptimisticModelId] =
    useOptimistic(selectedModelId);

  // 获取动态配置的模型
  const { data: dynamicModelsData, error: modelsError } = useSWR<{
    models: DynamicModel[];
    defaultModelId: string | null;
  }>("/api/models", fetcher);

  const dynamicModels = dynamicModelsData?.models ?? [];
  const defaultDynamicModelId = dynamicModelsData?.defaultModelId;

  // 新增：手动添加 Auto 选项
  const autoModel: DynamicModel = {
    id: "auto",
    modelId: "auto",
    displayName: "✨ 智能自动 (Auto)",
    description: "AI 自动识别意图，为您选择最佳模型",
    supportsVision: true,
    supportsImageGen: true,
    isDefault: false,
    providerName: "System",
    providerType: "system",
  };
  
  // 将 Auto 放在列表最前面
  const displayModels = [autoModel, ...dynamicModels];

  // 调试日志
  useEffect(() => {
    if (modelsError) {
      console.error("[ModelSelector] Failed to load models:", modelsError);
    }
    if (dynamicModelsData) {
      console.log("[ModelSelector] Loaded models:", dynamicModelsData);
    }
  }, [dynamicModelsData, modelsError]);


  // 如果当前选中的模型 ID 不存在，且有动态默认模型，则切换到默认模型
  useEffect(() => {
    if (dynamicModels.length > 0 && defaultDynamicModelId) {
      const currentExists = dynamicModels.some((m) => m.id === optimisticModelId);

      if (!currentExists) {
        startTransition(() => {
          setOptimisticModelId(defaultDynamicModelId);
          saveChatModelAsCookie(defaultDynamicModelId);
        });
      }
    }
  }, [
    dynamicModels,
    defaultDynamicModelId,
    optimisticModelId,
    setOptimisticModelId,
  ]);

  // 查找当前选中的模型（仅从数据库配置的模型中查找）
  const selectedModel = useMemo(() => {
    // 在动态模型中查找
    const dynamicModel = displayModels.find((m) => m.id === optimisticModelId);
    if (dynamicModel) {
      return {
        name: dynamicModel.displayName,
        description: dynamicModel.description,
      };
    }

    // 默认显示第一个可用模型
    if (displayModels.length > 0) {
      return {
        name: displayModels[0].displayName,
        description: displayModels[0].description,
      };
    }

    return { name: "未配置模型", description: null };
  }, [optimisticModelId, displayModels]);

  // 数据加载中
  if (!dynamicModelsData && !modelsError) {
    return (
      <Button variant="outline" disabled>
        加载中...
      </Button>
    );
  }

  // 数据加载完成，但没有配置任何模型
  if (dynamicModelsData && displayModels.length === 0) {
    return (
      <Button
        variant="outline"
        className="cursor-not-allowed opacity-50"
        disabled
      >
        ⚠️ 请在系统设置中配置 AI 模型
      </Button>
    );
  }

  // 数据加载失败
  if (modelsError) {
    return (
      <Button variant="outline" className="text-red-500" disabled>
        ⚠️ 加载失败
      </Button>
    );
  }

  return (
    <DropdownMenu onOpenChange={setOpen} open={open}>
      <DropdownMenuTrigger
        asChild
        className={cn(
          "w-fit data-[state=open]:bg-accent data-[state=open]:text-accent-foreground",
          className
        )}
      >
        <Button
          className="md:h-[34px] md:px-2 !ring-0 !ring-offset-0 !focus:ring-0 !focus-visible:ring-0 !outline-none"
          data-testid="model-selector"
          variant="outline"
        >
          {selectedModel.name}
          <ChevronDownIcon />
        </Button>
      </DropdownMenuTrigger>
      <DropdownMenuContent
        align="start"
        className="min-w-[280px] max-w-[90vw] sm:min-w-[300px]"
      >
        {/* 系统设置中配置的模型 */}
        {displayModels.map((model) => (
            <DropdownMenuItem
              asChild
                data-active={model.id === optimisticModelId}
                data-testid={`model-selector-item-${model.id}`}
                key={model.id}
              onSelect={() => {
                setOpen(false);

                startTransition(() => {
                    setOptimisticModelId(model.id);
                    saveChatModelAsCookie(model.id);
                });
              }}
            >
              <button
                className="group/item flex w-full flex-row items-center justify-between gap-2 sm:gap-4"
                type="button"
              >
                <div className="flex flex-col items-start gap-1">
                    <div className="flex items-center gap-2 text-sm sm:text-base">
                      {model.displayName}
                      {model.supportsVision && (
                        <span className="text-[10px] bg-green-100 text-green-700 px-1.5 py-0.5 rounded">
                          视觉
                        </span>
                      )}
                      {model.isDefault && (
                        <span className="text-[10px] bg-blue-100 text-blue-700 px-1.5 py-0.5 rounded">
                          默认
                        </span>
                      )}
                    </div>
                  <div className="line-clamp-2 text-muted-foreground text-xs">
                      {model.description ?? model.providerName}
                  </div>
                </div>

                <div className="shrink-0 text-foreground opacity-0 group-data-[active=true]/item:opacity-100 dark:text-foreground">
                  <CheckCircleFillIcon />
                </div>
              </button>
            </DropdownMenuItem>
            ))}
      </DropdownMenuContent>
    </DropdownMenu>
  );
}