import {  type ProviderDriverKind,  type ProviderInstanceId,  type ProviderOptionDescriptor,  type ProviderOptionSelection,  type ScopedThreadRef,  type ServerProviderModel,} from "@t3tools/contracts";import {  applyClaudePromptEffortPrefix,  buildProviderOptionSelectionsFromDescriptors,  getProviderOptionCurrentLabel,  getProviderOptionCurrentValue,  getProviderOptionDescriptors,  isClaudeUltrathinkPrompt,} from "@t3tools/shared/model";import { memo, useCallback, useState } from "react";import type { VariantProps } from "class-variance-authority";import { ChevronDownIcon } from "lucide-react";import { Button, buttonVariants } from "../ui/button";import {  Menu,  MenuGroup,  MenuPopup,  MenuRadioGroup,  MenuRadioItem,  MenuSeparator as MenuDivider,  MenuTrigger,} from "../ui/menu";import { useComposerDraftStore, DraftId } from "../../composerDraftStore";import { getProviderModelCapabilities } from "../../providerModels";import { cn } from "~/lib/utils"; type ProviderOptions = ReadonlyArray<ProviderOptionSelection>; type TraitsPersistence =  | {      threadRef?: ScopedThreadRef;      draftId?: DraftId;      onModelOptionsChange?: never;    }  | {      threadRef?: undefined;      onModelOptionsChange: (nextOptions: ProviderOptions | undefined) => void;    }; const ULTRATHINK_PROMPT_PREFIX = "Ultrathink:\n"; function replaceDescriptorCurrentValue(  descriptors: ReadonlyArray<ProviderOptionDescriptor>,  descriptorId: string,  currentValue: string | boolean | undefined,): ReadonlyArray<ProviderOptionDescriptor> {  return descriptors.map((descriptor) =>    descriptor.id !== descriptorId      ? descriptor      : descriptor.type === "boolean"        ? {            ...descriptor,            ...(typeof currentValue === "boolean" ? { currentValue } : {}),          }        : {            ...descriptor,            ...(typeof currentValue === "string" ? { currentValue } : {}),          },  );} function getDescriptorStringValue(  descriptor: Extract<ProviderOptionDescriptor, { type: "select" }> | null,): string | null {  if (!descriptor) {    return null;  }  const value = getProviderOptionCurrentValue(descriptor);  return typeof value === "string" ? value : null;} function getSelectedTraits(  provider: ProviderDriverKind,  models: ReadonlyArray<ServerProviderModel>,  model: string | null | undefined,  prompt: string,  modelOptions: ProviderOptions | null | undefined,  allowPromptInjectedEffort: boolean,) {  const caps = getProviderModelCapabilities(models, model, provider);  const descriptors = getProviderOptionDescriptors({    caps,    selections: modelOptions,  });  const selectDescriptors = descriptors.filter(    (descriptor): descriptor is Extract<ProviderOptionDescriptor, { type: "select" }> =>      descriptor.type === "select",  );  const booleanDescriptors = descriptors.filter(    (descriptor): descriptor is Extract<ProviderOptionDescriptor, { type: "boolean" }> =>      descriptor.type === "boolean",  );  const primarySelectDescriptor = selectDescriptors[0] ?? null;  const contextWindowDescriptor =    selectDescriptors.find((descriptor) => descriptor.id === "contextWindow") ?? null;  const agentDescriptor = selectDescriptors.find((descriptor) => descriptor.id === "agent") ?? null;  const fastModeDescriptor =    booleanDescriptors.find((descriptor) => descriptor.id === "fastMode") ?? null;  const thinkingDescriptor =    booleanDescriptors.find((descriptor) => descriptor.id === "thinking") ?? null;   // Prompt-controlled effort (e.g. ultrathink in prompt text)  const ultrathinkPromptControlled =    allowPromptInjectedEffort &&    (primarySelectDescriptor?.promptInjectedValues?.length ?? 0) > 0 &&    isClaudeUltrathinkPrompt(prompt);   // Check if "ultrathink" appears in the body text (not just our prefix)  const ultrathinkInBodyText =    ultrathinkPromptControlled && isClaudeUltrathinkPrompt(prompt.replace(/^Ultrathink:\s*/i, ""));  const effort =    (ultrathinkPromptControlled      ? "ultrathink"      : getDescriptorStringValue(primarySelectDescriptor)) ?? null;  const thinkingEnabled =    typeof thinkingDescriptor?.currentValue === "boolean" ? thinkingDescriptor.currentValue : null;  const fastModeEnabled =    typeof fastModeDescriptor?.currentValue === "boolean" ? fastModeDescriptor.currentValue : false;  const contextWindow = getDescriptorStringValue(contextWindowDescriptor);  const selectedAgent = getDescriptorStringValue(agentDescriptor);  const selectedAgentLabel = agentDescriptor    ? getProviderOptionCurrentLabel(agentDescriptor)    : null;   return {    caps,    descriptors,    selectDescriptors,    booleanDescriptors,    primarySelectDescriptor,    contextWindowDescriptor,    agentDescriptor,    fastModeDescriptor,    thinkingDescriptor,    effort,    thinkingEnabled,    fastModeEnabled,    contextWindow,    ultrathinkPromptControlled,    ultrathinkInBodyText,    selectedAgent,    selectedAgentLabel,  };} function getTraitsSectionVisibility(input: {  provider: ProviderDriverKind;  models: ReadonlyArray<ServerProviderModel>;  model: string | null | undefined;  prompt: string;  modelOptions: ProviderOptions | null | undefined;  allowPromptInjectedEffort?: boolean;}) {  const selected = getSelectedTraits(    input.provider,    input.models,    input.model,    input.prompt,    input.modelOptions,    input.allowPromptInjectedEffort ?? true,  );   const showEffort = selected.primarySelectDescriptor !== null;  const showThinking = selected.thinkingDescriptor !== null;  const showFastMode = selected.fastModeDescriptor !== null;  const showContextWindow = selected.contextWindowDescriptor !== null;  const showAgent = selected.agentDescriptor !== null;   return {    ...selected,    showEffort,    showThinking,    showFastMode,    showContextWindow,    showAgent,    hasAnyControls: showEffort || showThinking || showFastMode || showContextWindow || showAgent,  };} export function shouldRenderTraitsControls(input: {  provider: ProviderDriverKind;  models: ReadonlyArray<ServerProviderModel>;  model: string | null | undefined;  prompt: string;  modelOptions: ProviderOptions | null | undefined;  allowPromptInjectedEffort?: boolean;}): boolean {  return getTraitsSectionVisibility(input).hasAnyControls;} export interface TraitsMenuContentProps {  provider: ProviderDriverKind;  instanceId?: ProviderInstanceId;  models: ReadonlyArray<ServerProviderModel>;  model: string | null | undefined;  prompt: string;  onPromptChange: (prompt: string) => void;  modelOptions?: ProviderOptions | null | undefined;  allowPromptInjectedEffort?: boolean;  triggerVariant?: VariantProps<typeof buttonVariants>["variant"];  triggerClassName?: string;} export const TraitsMenuContent = memo(function TraitsMenuContentImpl({  provider,  instanceId,  models,  model,  prompt,  onPromptChange,  modelOptions,  allowPromptInjectedEffort = true,  ...persistence}: TraitsMenuContentProps & TraitsPersistence) {  const setProviderModelOptions = useComposerDraftStore((store) => store.setProviderModelOptions);  const updateModelOptions = useCallback(    (nextOptions: ProviderOptions | undefined) => {      if ("onModelOptionsChange" in persistence) {        persistence.onModelOptionsChange(nextOptions);        return;      }      const threadTarget = persistence.threadRef ?? persistence.draftId;      if (!threadTarget) {        return;      }      setProviderModelOptions(threadTarget, provider, nextOptions, {        ...(instanceId ? { instanceId } : {}),        model,        persistSticky: true,      });    },    [instanceId, model, persistence, provider, setProviderModelOptions],  );  const {    descriptors,    selectDescriptors,    booleanDescriptors,    primarySelectDescriptor,    ultrathinkPromptControlled,    ultrathinkInBodyText,    hasAnyControls,  } = getTraitsSectionVisibility({    provider,    models,    model,    prompt,    modelOptions,    allowPromptInjectedEffort,  });  const updateDescriptors = (nextDescriptors: ReadonlyArray<ProviderOptionDescriptor>) => {    updateModelOptions(buildProviderOptionSelectionsFromDescriptors(nextDescriptors));  };   const handleSelectChange = (    descriptor: Extract<ProviderOptionDescriptor, { type: "select" }>,    value: string,  ) => {    if (!value) return;    if (descriptor.promptInjectedValues?.includes(value)) {      const nextPrompt =        prompt.trim().length === 0          ? ULTRATHINK_PROMPT_PREFIX          : applyClaudePromptEffortPrefix(prompt, "ultrathink");      onPromptChange(nextPrompt);      return;    }    if (ultrathinkInBodyText && descriptor.id === primarySelectDescriptor?.id) return;    if (ultrathinkPromptControlled && descriptor.id === primarySelectDescriptor?.id) {      const stripped = prompt.replace(/^Ultrathink:\s*/i, "");      onPromptChange(stripped);    }    updateDescriptors(replaceDescriptorCurrentValue(descriptors, descriptor.id, value));  };   if (!hasAnyControls) {    return null;  }   return (    <>      {selectDescriptors.map((descriptor, index) => (        <div key={descriptor.id}>          {index > 0 ? <MenuDivider /> : null}          <MenuGroup>            <div className="px-2 pt-1.5 pb-1 font-medium text-muted-foreground text-xs">              {descriptor.label}            </div>            {ultrathinkInBodyText && descriptor.id === primarySelectDescriptor?.id ? (              <div className="px-2 pb-1.5 text-muted-foreground/80 text-xs">                Your prompt contains &quot;ultrathink&quot; in the text. Remove it to change this                option.              </div>            ) : null}            <MenuRadioGroup              value={                ultrathinkPromptControlled && descriptor.id === primarySelectDescriptor?.id                  ? "ultrathink"                  : (getDescriptorStringValue(descriptor) ?? "")              }              onValueChange={(value) => handleSelectChange(descriptor, value)}            >              {descriptor.options.map((option) => (                <MenuRadioItem                  key={option.id}                  value={option.id}                  disabled={ultrathinkInBodyText && descriptor.id === primarySelectDescriptor?.id}                >                  {option.label}                  {option.isDefault ? " (default)" : ""}                </MenuRadioItem>              ))}            </MenuRadioGroup>          </MenuGroup>        </div>      ))}      {booleanDescriptors.map((descriptor, index) => (        <div key={descriptor.id}>          {index > 0 || selectDescriptors.length > 0 ? <MenuDivider /> : null}          <MenuGroup>            <div className="px-2 py-1.5 font-medium text-muted-foreground text-xs">              {descriptor.label}            </div>            <MenuRadioGroup              value={descriptor.currentValue === true ? "on" : "off"}              onValueChange={(value) => {                updateDescriptors(                  replaceDescriptorCurrentValue(descriptors, descriptor.id, value === "on"),                );              }}            >              <MenuRadioItem value="on">On</MenuRadioItem>              <MenuRadioItem value="off">Off</MenuRadioItem>            </MenuRadioGroup>          </MenuGroup>        </div>      ))}    </>  );}); export const TraitsPicker = memo(function TraitsPicker({  provider,  instanceId,  models,  model,  prompt,  onPromptChange,  modelOptions,  allowPromptInjectedEffort = true,  triggerVariant,  triggerClassName,  ...persistence}: TraitsMenuContentProps & TraitsPersistence) {  const [isMenuOpen, setIsMenuOpen] = useState(false);  const { descriptors, primarySelectDescriptor, ultrathinkPromptControlled } =    getTraitsSectionVisibility({      provider,      models,      model,      prompt,      modelOptions,      allowPromptInjectedEffort,    });  if (    !shouldRenderTraitsControls({      provider,      models,      model,      prompt,      modelOptions,      allowPromptInjectedEffort,    })  ) {    return null;  }   const triggerLabels: Array<string> = [];  for (const descriptor of descriptors) {    const label =      ultrathinkPromptControlled && descriptor.id === primarySelectDescriptor?.id        ? "Ultrathink"        : descriptor.type === "boolean"          ? descriptor.id === "fastMode"            ? descriptor.currentValue === true              ? "Fast"              : "Normal"            : `${descriptor.label} ${descriptor.currentValue === true ? "On" : "Off"}`          : getProviderOptionCurrentLabel(descriptor);    if (typeof label === "string" && label.length > 0) {      triggerLabels.push(label);    }  }  const triggerLabel = triggerLabels.join(" · ");   const isCodexStyle = provider === "codex";   return (    <Menu      open={isMenuOpen}      onOpenChange={(open) => {        setIsMenuOpen(open);      }}    >      <MenuTrigger        render={          <Button            size="sm"            variant={triggerVariant ?? "ghost"}            className={cn(              isCodexStyle                ? "min-w-0 max-w-40 shrink justify-start overflow-hidden whitespace-nowrap px-2 text-muted-foreground/70 hover:text-foreground/80 sm:max-w-48 sm:px-3 [&_svg]:mx-0"                : "shrink-0 whitespace-nowrap px-2 text-muted-foreground/70 hover:text-foreground/80 sm:px-3",              triggerClassName,            )}          />        }      >        {isCodexStyle ? (          <span className="flex min-w-0 w-full items-center gap-2 overflow-hidden">            {triggerLabel}            <ChevronDownIcon aria-hidden="true" className="size-3 shrink-0 opacity-60" />          </span>        ) : (          <>            <span>{triggerLabel}</span>            <ChevronDownIcon aria-hidden="true" className="size-3 opacity-60" />          </>        )}      </MenuTrigger>      <MenuPopup align="start">        <TraitsMenuContent          provider={provider}          {...(instanceId ? { instanceId } : {})}          models={models}          model={model}          prompt={prompt}          onPromptChange={onPromptChange}          modelOptions={modelOptions}          allowPromptInjectedEffort={allowPromptInjectedEffort}          {...persistence}        />      </MenuPopup>    </Menu>  );});