import type { EmbeddingProvider } from "@rakazo/adapter-kit";
import {
  assertAllowedOpenAiCompatibleUrl,
  assertHttpsForKeyedOpenAiCompatibleUrl,
  OPENAI_COMPATIBLE_PROVIDER_ID,
} from "./openai-compatible-url.js";
import { createOpenAiCompatibleFetch } from "./pi-openai-compatible-provider.js";
import { readBodyCapped } from "./web-ssrf.js";

/** Texts sent per request; small enough for llama.cpp's default batch on short memories. */
export const OPENAI_COMPATIBLE_EMBEDDING_BATCH_SIZE = 32;
const EMBEDDING_TIMEOUT_MS = 30_000;
const MAX_EMBEDDING_RESPONSE_BYTES = 16 * 1024 * 1024;
const MAX_EMBEDDING_DIMENSIONS = 8_192;

export interface OpenAiCompatibleEmbeddingConfig {
  baseUrl: string;
  model: string;
  apiKey?: string;
}

/**
 * Embeddings from any server that speaks the OpenAI `POST /v1/embeddings` protocol, such as
 * llama.cpp (`llama-server --embeddings`), vLLM, Ollama, or a hosted API. Requests go through
 * the same address-checked fetch as OpenAI-compatible chat models.
 */
export class OpenAiCompatibleEmbeddingProvider implements EmbeddingProvider {
  private readonly endpoint: string;
  private readonly model: string;
  private readonly apiKey: string | undefined;
  private readonly fetch: typeof globalThis.fetch;

  constructor(config: OpenAiCompatibleEmbeddingConfig, fetchImpl?: typeof globalThis.fetch) {
    const baseUrl = assertAllowedOpenAiCompatibleUrl(config.baseUrl);
    this.apiKey = config.apiKey?.trim() || undefined;
    assertHttpsForKeyedOpenAiCompatibleUrl(baseUrl, this.apiKey);
    this.model = config.model.trim();
    if (!this.model) throw new Error("Embedding model is required");
    this.endpoint = new URL("embeddings", `${baseUrl.href}/`).href;
    this.fetch = createOpenAiCompatibleFetch(fetchImpl);
  }

  describe() {
    return {
      id: OPENAI_COMPATIBLE_PROVIDER_ID,
      contractVersion: "1",
      adapterVersion: "0.1.0",
      capabilities: { model: this.model },
    };
  }

  async embed(texts: string[], signal?: AbortSignal): Promise<number[][]> {
    const vectors: number[][] = [];
    for (let start = 0; start < texts.length; start += OPENAI_COMPATIBLE_EMBEDDING_BATCH_SIZE) {
      const batch = texts.slice(start, start + OPENAI_COMPATIBLE_EMBEDDING_BATCH_SIZE);
      vectors.push(...(await this.embedBatch(batch, signal)));
    }
    if (vectors.some((vector) => vector.length !== vectors[0]!.length)) {
      throw new Error("Embeddings server returned vectors of different lengths");
    }
    return vectors;
  }

  private async embedBatch(texts: string[], signal?: AbortSignal): Promise<number[][]> {
    const timeout = AbortSignal.timeout(EMBEDDING_TIMEOUT_MS);
    const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;
    const headers: Record<string, string> = {
      Accept: "application/json",
      "Content-Type": "application/json",
    };
    if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`;
    try {
      const response = await this.fetch(this.endpoint, {
        method: "POST",
        headers,
        body: JSON.stringify({ model: this.model, input: texts, encoding_format: "float" }),
        redirect: "error",
        signal: requestSignal,
      });
      if (!response.ok) {
        await response.body?.cancel().catch(() => undefined);
        throw new Error(`Embeddings server returned ${response.status}`);
      }
      const bytes = await readBodyCapped(response, MAX_EMBEDDING_RESPONSE_BYTES, requestSignal);
      let body: unknown;
      try {
        body = JSON.parse(new TextDecoder().decode(bytes));
      } catch {
        throw new Error("Embeddings server returned invalid JSON");
      }
      return parseEmbeddingResponse(body, texts.length);
    } catch (error) {
      if (signal?.aborted) throw error;
      if (timeout.aborted) throw new Error("Embeddings server timed out");
      // undici reports refused or reset connections as a bare "fetch failed".
      if (error instanceof TypeError && error.message === "fetch failed") {
        throw new Error("Embeddings server is unreachable");
      }
      throw error;
    }
  }
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
}

function isVector(value: unknown): value is number[] {
  return (
    Array.isArray(value) &&
    value.length > 0 &&
    value.length <= MAX_EMBEDDING_DIMENSIONS &&
    value.every((entry) => typeof entry === "number" && Number.isFinite(entry))
  );
}

/** Reads `{data: [{index, embedding}]}`, honoring `index` when the server reorders entries. */
export function parseEmbeddingResponse(body: unknown, count: number): number[][] {
  const data = isRecord(body) && Array.isArray(body.data) ? body.data : null;
  if (!data || data.length !== count) {
    throw new Error("Embeddings server returned an unexpected number of vectors");
  }
  const vectors: Array<number[] | undefined> = new Array(count);
  data.forEach((entry, position) => {
    const index = isRecord(entry) && Number.isInteger(entry.index) ? Number(entry.index) : position;
    const embedding = isRecord(entry) ? entry.embedding : undefined;
    if (!isVector(embedding) || index < 0 || index >= count || vectors[index]) {
      throw new Error("Embeddings server returned a malformed vector");
    }
    vectors[index] = embedding;
  });
  return vectors as number[][];
}