ModelPickerGroupsBuilder.java

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

package com.microsoft.copilot.eclipse.ui.chat;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

import org.apache.commons.lang3.StringUtils;
import org.eclipse.swt.SWT;
import org.eclipse.swt.graphics.Image;
import org.eclipse.swt.widgets.Display;

import com.microsoft.copilot.eclipse.core.lsp.protocol.CopilotModel;
import com.microsoft.copilot.eclipse.ui.i18n.Messages;
import com.microsoft.copilot.eclipse.ui.preferences.ByokPreferencePage;
import com.microsoft.copilot.eclipse.ui.swt.DropdownItem;
import com.microsoft.copilot.eclipse.ui.swt.DropdownItemGroup;
import com.microsoft.copilot.eclipse.ui.swt.ModelHoverContentProvider;
import com.microsoft.copilot.eclipse.ui.utils.ModelUtils;
import com.microsoft.copilot.eclipse.ui.utils.PreferencesUtils;
import com.microsoft.copilot.eclipse.ui.utils.SwtUtils;
import com.microsoft.copilot.eclipse.ui.utils.UiUtils;

/**
 * Builds model picker dropdown groups for the chat UI.
 */
public final class ModelPickerGroupsBuilder {

  private static Image warningIcon;

  private ModelPickerGroupsBuilder() {
  }

  /**
   * Builds grouped dropdown items for the model picker.
   *
   * @param modelMap available models keyed by id
   * @param showAddPremiumModelOption whether to include the premium upsell action
   * @param showByokManageOption whether to include the BYOK manage action
   * @return grouped dropdown items for the model picker
   */
  public static List<DropdownItemGroup> build(Map<String, CopilotModel> modelMap, boolean showAddPremiumModelOption,
      boolean showByokManageOption) {
    return build(modelMap, showAddPremiumModelOption, showByokManageOption, null);
  }

  /**
   * Builds grouped dropdown items for the model picker, including the effective reasoning effort in the suffix.
   *
   * @param modelMap available models keyed by id
   * @param showAddPremiumModelOption whether to include the premium upsell action
   * @param showByokManageOption whether to include the BYOK manage action
   * @param reasoningEffortResolver resolves the effective reasoning effort for a given model (user-selected when
   *     present, otherwise the inferred default), or {@code null} when none applies
   * @return grouped dropdown items for the model picker
   */
  public static List<DropdownItemGroup> build(Map<String, CopilotModel> modelMap, boolean showAddPremiumModelOption,
      boolean showByokManageOption, Function<CopilotModel, String> reasoningEffortResolver) {
    List<CopilotModel> otherModels = new ArrayList<>();
    List<CopilotModel> standardModels = new ArrayList<>();
    List<CopilotModel> premiumModels = new ArrayList<>();
    List<CopilotModel> customModels = new ArrayList<>();

    for (CopilotModel model : modelMap.values()) {
      if (model.getProviderName() != null) {
        customModels.add(model);
      } else if (model.getBilling() != null) {
        if (model.getBilling().isPremium()) {
          premiumModels.add(model);
        } else {
          standardModels.add(model);
        }
      } else {
        otherModels.add(model);
      }
    }

    standardModels.sort((a, b) -> String.CASE_INSENSITIVE_ORDER.compare(a.getModelName(), b.getModelName()));
    premiumModels.sort((a, b) -> String.CASE_INSENSITIVE_ORDER.compare(a.getModelName(), b.getModelName()));
    customModels.sort((a, b) -> String.CASE_INSENSITIVE_ORDER.compare(a.getModelName(), b.getModelName()));

    List<DropdownItemGroup> groups = new ArrayList<>();
    if (!otherModels.isEmpty()) {
      groups.add(DropdownItemGroup.of(buildModelDropdownItems(otherModels, reasoningEffortResolver)));
    }
    if (!standardModels.isEmpty()) {
      groups.add(DropdownItemGroup.of(Messages.chat_standardModels,
          buildModelDropdownItems(standardModels, reasoningEffortResolver)));
    }
    if (!premiumModels.isEmpty()) {
      String header = standardModels.isEmpty() ? Messages.chat_copilotModels : Messages.chat_premiumModels;
      groups.add(DropdownItemGroup.of(header, buildModelDropdownItems(premiumModels, reasoningEffortResolver)));
    }
    if (!customModels.isEmpty()) {
      groups.add(DropdownItemGroup.of(Messages.chat_customModels,
          buildModelDropdownItems(customModels, reasoningEffortResolver)));
    }

    List<DropdownItem> actionItems = new ArrayList<>();
    if (showAddPremiumModelOption) {
      actionItems.add(new DropdownItem.Builder().label(Messages.chat_addPremiumModels).onAction(
          () -> UiUtils.executeCommandWithParameters("com.microsoft.copilot.eclipse.commands.upgradeCopilotPlan", null))
          .build());
    }
    if (showByokManageOption) {
      actionItems.add(new DropdownItem.Builder().label(Messages.chat_actionBar_modelPicker_manageModels)
          .onAction(ModelPickerGroupsBuilder::openManageModelsPreferences).build());
    }
    if (!actionItems.isEmpty()) {
      groups.add(DropdownItemGroup.of(actionItems));
    }
    return groups;
  }

  private static List<DropdownItem> buildModelDropdownItems(List<CopilotModel> models,
      Function<CopilotModel, String> reasoningEffortResolver) {
    List<DropdownItem> items = new ArrayList<>();
    for (CopilotModel model : models) {
      String rawName = model.getModelName();
      boolean alreadyHasPreview = rawName != null && rawName.toLowerCase().endsWith("(preview)");
      String name = model.isPreview() && !alreadyHasPreview ? rawName + " " + Messages.model_preview_suffix : rawName;

      String effectiveEffort = reasoningEffortResolver != null ? reasoningEffortResolver.apply(model) : null;
      String suffix = ModelUtils.getModelSuffix(model, effectiveEffort);
      String effortLevel = ModelUtils.formatReasoningEffortLevel(effectiveEffort);
      String selectedLabel = StringUtils.isNotBlank(effortLevel) && StringUtils.isNotBlank(name)
          ? name + " - " + effortLevel : null;

      items.add(new DropdownItem.Builder().id(rawName).label(name).selectedLabel(selectedLabel).suffix(suffix)
          .icon(resolveModelIcon(model)).hoverProvider(new ModelHoverContentProvider(model)).build());
    }
    return items;
  }

  private static Image resolveModelIcon(CopilotModel model) {
    if (StringUtils.isBlank(model.getDegradationReason())) {
      return null;
    }
    if (warningIcon == null || warningIcon.isDisposed()) {
      warningIcon = UiUtils.isDarkTheme()
          ? UiUtils.buildImageFromPngPath("/icons/dropdown/dropdown_warning_dark.png")
          : UiUtils.buildImageFromPngPath("/icons/dropdown/dropdown_warning.png");
      Display display = SwtUtils.getDisplay();
      if (display != null) {
        display.addListener(SWT.Dispose, event -> disposeWarningIcons());
      }
    }
    return warningIcon;
  }

  private static void disposeWarningIcons() {
    if (warningIcon != null && !warningIcon.isDisposed()) {
      warningIcon.dispose();
      warningIcon = null;
    }
  }

  private static void openManageModelsPreferences() {
    Map<String, Object> parameters = new HashMap<>();
    parameters.put("com.microsoft.copilot.eclipse.commands.openPreferences.activePageId", ByokPreferencePage.ID);
    parameters.put("com.microsoft.copilot.eclipse.commands.openPreferences.pageIds",
        String.join(",", PreferencesUtils.getAllPreferenceIds()));
    UiUtils.executeCommandWithParameters("com.microsoft.copilot.eclipse.commands.openPreferences", parameters);
  }
}