import { firstValueFrom } from 'rxjs';
import { LLMDescriptor, ModelParameters } from '@models/flow';
import { FieldRetriever } from '@services/retriever/field-retriever';
import { NodeSettingField, NodeSettingsDialogService } from '@services/dialogs/node-settings-dialog';
import { readSimulatorParameters } from '@shared/task-execution-viewer/execution-viewer.utils';

/** The retrievers the editor already uses for a node's provider and model fields. */
const PROVIDER_RETRIEVER_URL = '/retriever/LLM/providers';
const MODEL_RETRIEVER_URL = '/retriever/LLM/models';

const PARAMETER_GROUP = 'Model parameters';

/**
 * The optional sampling knobs, behind a section that starts closed. Provider and model are what
 * anyone opening this dialog came for; these are for the runs where you already know you want
 * them, and shown flat they made the common case look like a five-field form.
 *
 * <p>These are the same five fields a node's own parameter panel shows, and they are declared to
 * behave the same way: `defaultsWhenEmpty` so an empty box reads as "the provider decides" and
 * there is a way back to it after typing, the bounds the server actually enforces, and an arrow
 * step a decimal can move by. Hand-written rather than derived from the published schema the way
 * `buildSchemaObjectDialog` does it: this dialog is picking a model for a run, not editing a node's
 * configuration, and reaching for a block type's schema to render five known fields would buy a
 * network call and a failure mode. The labels, tips and bounds are ModelParameters' own - keep them
 * in step with it.
 */
const PARAMETER_FIELDS: NodeSettingField[] = [
  {
    key: 'temperature', label: 'Temperature', type: 'number', group: PARAMETER_GROUP,
    // The server caps this at 1.0; offering 2 here only produced a value it would reject.
    min: 0, max: 1, stepIncrement: 0.1,
    defaultsWhenEmpty: true, placeholder: 'Leave empty for the default',
    tip: 'Higher values make the output more varied. 0 makes it as repeatable as the model allows.'
  },
  {
    key: 'topP', label: 'Top P', type: 'number', group: PARAMETER_GROUP,
    min: 0, max: 1, stepIncrement: 0.1,
    defaultsWhenEmpty: true, placeholder: 'Leave empty for the default',
    tip: 'Nucleus sampling: consider only the most likely tokens adding up to this probability.'
  },
  {
    key: 'topK', label: 'Top K', type: 'number', group: PARAMETER_GROUP,
    min: 1, step: 1, stepIncrement: 1,
    defaultsWhenEmpty: true, placeholder: 'Leave empty for the default',
    tip: 'Consider only this many candidate tokens at each step.'
  },
  {
    key: 'maxTokens', label: 'Max tokens', type: 'number', group: PARAMETER_GROUP,
    min: 1, step: 1, stepIncrement: 1,
    defaultsWhenEmpty: true, placeholder: 'Leave empty for the default',
    tip: 'Upper bound on the length of the generated answer.'
  },
  {
    key: 'seed', label: 'Seed', type: 'number', group: PARAMETER_GROUP,
    step: 1, stepIncrement: 1,
    defaultsWhenEmpty: true, placeholder: 'Leave empty for the default',
    tip: 'Fixes the randomness, so the same inputs give the same answer. Needed to tell a real change from model noise.'
  }
];

export type LLMDescriptorSettingsRequest = {
  title: string;
  /** Sampling defaults offered in the closed section, e.g. temperature 0 for a judge. */
  defaultParameters?: Partial<Record<keyof ModelParameters, number>>;
  /**
   * What to start from, when the caller already has a model worth repeating - the simulator of the
   * run this one is a rerun of. Preselecting it is the difference between a comparison of one
   * intervention and a comparison of two different simulators.
   */
  initialDescriptor?: LLMDescriptor | null;
};

/**
 * Asks for a provider, a model and optionally how to sample from them.
 *
 * <p>One dialog for every place that picks a model at the moment it is needed rather than in a
 * node's configuration: the interaction simulator, and the LLM assessment of a bias comparison.
 * They ask the same question, so they ask it with the same window - the provider list, the models
 * that follow from the chosen provider, and the same optional parameters.
 *
 * Returns null when the dialog is dismissed. Throws when no provider is published at all: the
 * caller has a place to say so, and answering null there made the button that opened it look broken.
 */
export async function openLLMDescriptorSettings(
  settingsDialog: NodeSettingsDialogService,
  fieldRetriever: FieldRetriever,
  request: LLMDescriptorSettingsRequest
): Promise<LLMDescriptor | null> {
  const providerOptions = await loadOptions(fieldRetriever, 'providers', {}, PROVIDER_RETRIEVER_URL);
  if (!providerOptions.length) {
    throw new Error('No LLM provider is published by the server, so there is no model to choose.');
  }

  const inherited = request.initialDescriptor ?? null;
  // Only if it is still on offer: a provider that has since been unregistered would leave the
  // dialog pointing at a model nobody can pick.
  const inheritedProvider = inherited?.provider
    && providerOptions.some((option) => option.value === inherited.provider)
    ? inherited.provider
    : null;

  const defaultProvider = inheritedProvider ?? providerOptions[0].value;
  const initialModelOptions = await loadOptions(
    fieldRetriever,
    'models',
    { provider: defaultProvider },
    MODEL_RETRIEVER_URL
  );

  const inheritedParameters: Record<string, string> = Object.fromEntries(
    Object.entries(inheritedProvider ? inherited?.parameters ?? {} : {})
      .filter(([, value]) => value !== null && value !== undefined)
      .map(([key, value]) => [key, String(value)])
  );

  const buildFields = (
    providers: { label: string; value: string; }[],
    models: { label: string; value: string; }[]
  ): NodeSettingField[] => [
    { key: 'provider', label: 'Provider', type: 'select', options: providers, required: true, autofocus: true },
    { key: 'model', label: 'Model', type: 'select', options: models, required: true },
    // Inherited sampling is opened, not hidden: a seed carried over from the run being repeated is
    // the reason the two runs are comparable, and behind a closed section nobody would see it.
    ...PARAMETER_FIELDS.map((field) => (inheritedParameters[field.key] === undefined
      ? field
      : { ...field, group: undefined }))
  ];

  const defaults = Object.fromEntries(
    Object.entries(request.defaultParameters ?? {}).map(([key, value]) => [key, String(value)])
  );

  const inheritedModel = inheritedProvider && inherited?.model
    && initialModelOptions.some((option) => option.value === inherited.model)
    ? inherited.model
    : null;

  const result = await settingsDialog.open({
    title: request.title,
    fields: buildFields(providerOptions, initialModelOptions),
    initial: {
      provider: defaultProvider,
      model: inheritedModel ?? initialModelOptions[0]?.value ?? '',
      ...defaults,
      // Last: what the run being repeated used wins over a generic default.
      ...inheritedParameters
    },
    onValuesChange: async (draft) => {
      const provider = String(draft['provider'] ?? '').trim();
      const modelOptions = provider
        ? await loadOptions(fieldRetriever, 'models', { provider }, MODEL_RETRIEVER_URL)
        : [];

      return {
        fields: buildFields(providerOptions, modelOptions),
        initial: {
          provider,
          model: modelOptions[0]?.value ?? ''
        }
      };
    }
  });

  if (!result) return null;

  const provider = String(result['provider'] ?? '').trim();
  const model = String(result['model'] ?? '').trim();
  if (!provider || !model) {
    return null;
  }

  const parameters = readSimulatorParameters(result);
  return parameters ? { provider, model, parameters } : { provider, model };
}

async function loadOptions(
  fieldRetriever: FieldRetriever,
  key: string,
  context: Record<string, string>,
  retrieverUrl: string
): Promise<Array<{ label: string; value: string }>> {
  const values = await firstValueFrom(fieldRetriever.retrieveValues('LLM', key, context, retrieverUrl));
  return values.map((value) => ({ label: value, value }));
}