import type { InferenceServer } from '@sirene/shared';
import { ServiceUnavailableError } from '../errors';
import type { InferenceTarget } from '../types';
import { inferenceServerService } from './inference-server.service';
import { serverModelsService } from './server-models.service';

class NoInferenceServerError extends ServiceUnavailableError {
  public constructor(message = 'No inference server is available') {
    super('model.noOnlineServer', message);
  }
}

const inFlight = new Map<string, number>();

interface PickOptions {
  /** Restrict candidates to servers that have this model installed. */
  requireModel?: string;
}

function targetOf(server: InferenceServer): InferenceTarget {
  return { url: server.url, authToken: server.authToken };
}

/** Pick a server for an outgoing call. Strategy:
 *  - Filter to enabled+online (or unknown if nothing online) servers
 *  - If `requireModel` is set, keep only servers that have it installed
 *  - Among candidates: pick the one with the fewest in-flight calls; ties broken by priority */
async function pickServer(options: PickOptions = {}): Promise<InferenceServer> {
  const all = await inferenceServerService.listEnabled();
  if (all.length === 0) {
    throw new NoInferenceServerError('No inference server is configured. Add one from Settings.');
  }

  let candidates = all.filter((s) => s.lastHealth.status === 'online');
  if (candidates.length === 0) {
    // Treat 'unknown' as a candidate so a freshly-added server can still be tried before
    // its first health probe completes.
    candidates = all.filter((s) => !s.lastHealth.status || s.lastHealth.status === 'unknown');
  }
  if (candidates.length === 0) {
    throw new NoInferenceServerError('All configured inference servers are offline.');
  }

  if (options.requireModel) {
    const serverIds = new Set(await serverModelsService.serversWithModel(options.requireModel));
    const filtered = candidates.filter((s) => serverIds.has(s.id));
    if (filtered.length === 0) {
      throw new NoInferenceServerError(`Model "${options.requireModel}" is not installed on any online server.`);
    }
    candidates = filtered;
  }

  candidates.sort((a, b) => {
    const aLoad = inFlight.get(a.id) ?? 0;
    const bLoad = inFlight.get(b.id) ?? 0;
    if (aLoad !== bLoad) {
      return aLoad - bLoad;
    }
    return b.priority - a.priority;
  });

  const winner = candidates[0];
  if (!winner) {
    throw new NoInferenceServerError();
  }
  return winner;
}

export async function pickTarget(options: PickOptions = {}): Promise<InferenceTarget> {
  return targetOf(await pickServer(options));
}
