import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const BASE_URL = "http://127.0.0.1:1234/v1";
const DEFAULT_CONTEXT_WINDOW = 128_000;
const DEFAULT_MAX_TOKENS = 16_384;

type LmStudioModel = {
  id?: unknown;
  name?: unknown;
  context_window?: unknown;
  max_tokens?: unknown;
};

type LmStudioModelsResponse = {
  data?: LmStudioModel[];
};

function positiveInteger(value: unknown, fallback: number): number {
  return typeof value === "number" && Number.isInteger(value) && value > 0
    ? value
    : fallback;
}

export default function (pi: ExtensionAPI) {
  pi.registerProvider("lm-studio", {
    name: "LM Studio",
    baseUrl: BASE_URL,
    apiKey: "lm-studio",
    api: "openai-completions",
    models: [],
    async refreshModels({ allowNetwork, signal, stored }) {
      if (!allowNetwork) return [...(stored?.models ?? [])];

      const response = await fetch(`${BASE_URL}/models`, { signal });
      if (!response.ok) {
        throw new Error(
          `LM Studio model discovery failed: ${response.status} ${response.statusText}`,
        );
      }

      const payload = (await response.json()) as LmStudioModelsResponse;
      if (!Array.isArray(payload.data)) {
        throw new Error("LM Studio returned an invalid /v1/models response");
      }

      return payload.data.flatMap((model) => {
        if (typeof model.id !== "string" || model.id.length === 0) return [];

        return [
          {
            id: model.id,
            name: typeof model.name === "string" ? model.name : model.id,
            reasoning: false,
            input: ["text"] as ["text"],
            cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
            contextWindow: positiveInteger(
              model.context_window,
              DEFAULT_CONTEXT_WINDOW,
            ),
            maxTokens: positiveInteger(model.max_tokens, DEFAULT_MAX_TOKENS),
            compat: {
              supportsDeveloperRole: false,
              supportsReasoningEffort: false,
              supportsUsageInStreaming: false,
              maxTokensField: "max_tokens" as const,
            },
          },
        ];
      });
    },
  });
}