import { platform, totalmem } from "node:os";

import type { NativeEvaluationOptions } from "../native-worker/evaluation";
/**
 * Embedding port implementation using node-llama-cpp.
 *
 * @module src/llm/nodeLlamaCpp/embedding
 */
import type { EmbeddingPort, LlmResult } from "../types";
import type { ModelManager } from "./lifecycle";

import { inferenceFailedError } from "../errors";
import { checkEvaluation, startEvaluation } from "../native-worker/evaluation";

// LlamaModel type from node-llama-cpp
type LlamaModel = Awaited<
  ReturnType<
    Awaited<ReturnType<typeof import("node-llama-cpp").getLlama>>["loadModel"]
  >
>;

type LlamaEmbeddingContext = Awaited<
  ReturnType<LlamaModel["createEmbeddingContext"]>
>;

type Llama = Awaited<ReturnType<typeof import("node-llama-cpp").getLlama>>;

interface EmbeddingWorker {
  context: LlamaEmbeddingContext;
  pending: number;
}

interface TokenizingModel {
  trainContextSize?: number;
  tokenize(text: string): readonly number[];
  detokenize(tokens: readonly number[]): string;
}

type EmbeddingInput = Parameters<LlamaEmbeddingContext["getEmbeddingFor"]>[0];

// Aim for a small pool so CPU-only runs can exploit parallel contexts without
// multiplying RAM usage too aggressively. Additional contexts fall back
// gracefully if memory is tight.
const MAX_DEFAULT_EMBEDDING_CONTEXTS = 2;
const MAX_EMBEDDING_CONTEXTS_OVERRIDE = 4;
const TARGET_CORES_PER_EMBEDDING_CONTEXT = 4;
const CONSTRAINED_WINDOWS_THRESHOLD_BYTES = 16 * 1024 * 1024 * 1024;
const MID_MEMORY_WINDOWS_THRESHOLD_BYTES = 24 * 1024 * 1024 * 1024;
const LOW_MEMORY_WINDOWS_CONTEXTS = 1;
const MID_MEMORY_WINDOWS_CONTEXTS = 2;
const DEFAULT_EMBEDDING_CONTEXT_SIZE = 2_048;

function embeddingVectorToArray(vector: readonly number[]): number[] {
  return Array.isArray(vector) ? (vector as number[]) : Array.from(vector);
}

function resolveEmbeddingContextPoolOverride(
  env: NodeJS.ProcessEnv = process.env
): number | undefined {
  const raw = env.GNO_EMBED_CONTEXTS;
  if (!raw) {
    return undefined;
  }
  const parsed = Number.parseInt(raw, 10);
  if (!(Number.isFinite(parsed) && parsed > 0)) {
    return undefined;
  }
  return Math.max(1, Math.min(MAX_EMBEDDING_CONTEXTS_OVERRIDE, parsed));
}

function resolveThreadsPerContextOverride(
  env: NodeJS.ProcessEnv = process.env
): number | undefined {
  const raw = env.GNO_EMBED_THREADS;
  if (!raw) {
    return undefined;
  }
  const parsed = Number.parseInt(raw, 10);
  if (!(Number.isFinite(parsed) && parsed > 0)) {
    return undefined;
  }
  return Math.max(1, parsed);
}

function resolveEmbeddingContextSizeOverride(
  env: NodeJS.ProcessEnv = process.env
): number | undefined {
  const raw = env.GNO_EMBED_CONTEXT_SIZE;
  if (!raw) {
    return undefined;
  }
  const parsed = Number.parseInt(raw, 10);
  if (!(Number.isFinite(parsed) && parsed > 0)) {
    return undefined;
  }
  return Math.max(128, parsed);
}

export function resolveEmbeddingContextPoolSize(options: {
  gpu: Llama["gpu"];
  cpuMathCores: number;
  env?: NodeJS.ProcessEnv;
  platformName?: NodeJS.Platform;
  totalMemoryBytes?: number;
}): number {
  if (options.gpu !== false) {
    return 1;
  }

  const override = resolveEmbeddingContextPoolOverride(options.env);
  if (override !== undefined) {
    return override;
  }

  const platformName = options.platformName ?? platform();
  const totalMemoryBytes = options.totalMemoryBytes ?? totalmem();
  if (
    platformName === "win32" &&
    totalMemoryBytes < CONSTRAINED_WINDOWS_THRESHOLD_BYTES
  ) {
    return LOW_MEMORY_WINDOWS_CONTEXTS;
  }

  const cpuMathCores = Math.max(1, options.cpuMathCores);
  const adaptivePoolSize = Math.max(
    1,
    Math.min(
      MAX_DEFAULT_EMBEDDING_CONTEXTS,
      Math.ceil(cpuMathCores / TARGET_CORES_PER_EMBEDDING_CONTEXT)
    )
  );

  if (
    platformName === "win32" &&
    totalMemoryBytes < MID_MEMORY_WINDOWS_THRESHOLD_BYTES
  ) {
    return Math.min(MID_MEMORY_WINDOWS_CONTEXTS, adaptivePoolSize);
  }

  return adaptivePoolSize;
}

export class NodeLlamaCppEmbedding implements EmbeddingPort {
  private workers: EmbeddingWorker[] = [];
  private contextsPromise: Promise<LlmResult<LlamaEmbeddingContext[]>> | null =
    null;
  private lifecycleVersion = 0;
  private dims: number | null = null;
  private llamaModel: TokenizingModel | null = null;
  private embeddingContextSize = DEFAULT_EMBEDDING_CONTEXT_SIZE;
  private warnedSingleTruncation = false;
  private warnedBatchTruncation = false;
  private readonly manager: ModelManager;
  readonly modelUri: string;
  private readonly modelPath: string;

  constructor(manager: ModelManager, modelUri: string, modelPath: string) {
    this.manager = manager;
    this.modelUri = modelUri;
    this.modelPath = modelPath;
  }

  async init(options?: NativeEvaluationOptions): Promise<LlmResult<void>> {
    const lease = this.manager.acquireLease(this.modelUri, false);
    try {
      checkEvaluation(options);
      const contexts = await this.getContexts(options);
      checkEvaluation(options);
      if (!contexts.ok) {
        return contexts;
      }
      return { ok: true, value: undefined };
    } finally {
      lease.release();
    }
  }

  async embed(
    text: string,
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<number[]>> {
    const lease = this.manager.acquireLease(this.modelUri);
    try {
      return await this.embedLeased(text, options);
    } finally {
      lease.release();
    }
  }

  private async embedLeased(
    text: string,
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<number[]>> {
    checkEvaluation(options);
    const contexts = await this.getContexts(options);
    checkEvaluation(options);
    if (!contexts.ok) {
      return contexts;
    }

    try {
      const prepared = this.truncateForEmbedding(text, "single");
      if (!prepared.ok) {
        return { ok: false, error: prepared.error };
      }
      startEvaluation(options);
      const embedding = await this.runOnWorker((worker) =>
        worker.context.getEmbeddingFor(prepared.value.input)
      );
      checkEvaluation(options);
      const vector = embeddingVectorToArray(embedding.vector);

      // Cache dimensions on first call
      if (this.dims === null) {
        this.dims = vector.length;
      }

      return { ok: true, value: vector };
    } catch (e) {
      return { ok: false, error: inferenceFailedError(this.modelUri, e) };
    }
  }

  async embedBatch(
    texts: string[],
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<number[][]>> {
    const lease = this.manager.acquireLease(this.modelUri);
    try {
      return await this.embedBatchLeased(texts, options);
    } finally {
      lease.release();
    }
  }

  private async embedBatchLeased(
    texts: string[],
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<number[][]>> {
    checkEvaluation(options);
    const contexts = await this.getContexts(options);
    checkEvaluation(options);
    if (!contexts.ok) {
      return contexts;
    }

    if (texts.length === 0) {
      return { ok: true, value: [] };
    }

    try {
      const preparedInputs: EmbeddingInput[] = [];
      for (const text of texts) {
        const prepared = this.truncateForEmbedding(text, "batch");
        if (!prepared.ok) {
          return { ok: false, error: prepared.error };
        }
        preparedInputs.push(prepared.value.input);
      }

      const allResults = Array.from(
        { length: texts.length },
        () => [] as number[]
      );
      let nextIndex = 0;

      const settled = await Promise.allSettled(
        this.workers.map(async (worker) => {
          while (true) {
            checkEvaluation(options);
            const index = nextIndex;
            nextIndex += 1;
            if (index >= preparedInputs.length) {
              return;
            }

            const input = preparedInputs[index];
            if (input === undefined) {
              return;
            }
            startEvaluation(options);
            const embedding = await this.runOnSpecificWorker(
              worker,
              (current) => current.context.getEmbeddingFor(input)
            );
            allResults[index] = embeddingVectorToArray(embedding.vector);
          }
        })
      );

      const firstRejection = settled.find(
        (result): result is PromiseRejectedResult =>
          result.status === "rejected"
      );
      if (firstRejection) {
        return {
          ok: false,
          error: inferenceFailedError(this.modelUri, firstRejection.reason),
        };
      }

      checkEvaluation(options);
      // Cache dimensions from first result
      const firstResult = allResults[0];
      if (this.dims === null && firstResult !== undefined) {
        this.dims = firstResult.length;
      }

      return { ok: true, value: allResults };
    } catch (e) {
      return { ok: false, error: inferenceFailedError(this.modelUri, e) };
    }
  }

  dimensions(): number {
    if (this.dims === null) {
      throw new Error("Call init() or embed() first to initialize dimensions");
    }
    return this.dims;
  }

  /** Child-only effective tokenizer policy; never infer this from parent config. */
  getContextIdentity(): {
    contextSize: number;
    truncationPolicy: string;
    contextCount: number;
    threadsPerContext: number;
  } {
    if (!this.llamaModel || !this.workers.length)
      throw new Error("Embedding context not initialized");
    const trainSize = this.llamaModel.trainContextSize;
    const limit =
      typeof trainSize === "number" &&
      Number.isFinite(trainSize) &&
      trainSize > 0
        ? Math.min(Math.floor(trainSize), this.embeddingContextSize)
        : this.embeddingContextSize;
    return {
      contextSize: this.embeddingContextSize,
      truncationPolicy: `truncate-tail-tokens-v1:limit=${Math.max(1, limit - 4)}`,
      contextCount: this.workers.length,
      threadsPerContext: this.threadsPerContext,
    };
  }

  async dispose(): Promise<void> {
    const lease = this.manager.acquireLease(this.modelUri, false);
    try {
      this.lifecycleVersion += 1;
      this.contextsPromise = null;
      this.llamaModel = null;
      this.dims = null;
      this.warnedSingleTruncation = false;
      this.warnedBatchTruncation = false;
      const workers = this.workers;
      this.workers = [];

      for (const worker of workers) {
        try {
          await worker.context.dispose();
        } catch {
          // Ignore disposal errors
        }
      }
    } finally {
      lease.release();
    }
  }

  private async runOnWorker<T>(
    task: (worker: EmbeddingWorker) => Promise<T>
  ): Promise<T> {
    const worker = this.getLeastBusyWorker();
    return this.runOnSpecificWorker(worker, task);
  }

  private async runOnSpecificWorker<T>(
    worker: EmbeddingWorker,
    task: (worker: EmbeddingWorker) => Promise<T>
  ): Promise<T> {
    worker.pending += 1;
    try {
      return await task(worker);
    } finally {
      worker.pending -= 1;
    }
  }

  private getLeastBusyWorker(): EmbeddingWorker {
    const firstWorker = this.workers[0];
    if (!firstWorker) {
      throw new Error("Embedding context not initialized");
    }

    let bestWorker = firstWorker;
    for (const worker of this.workers) {
      if (worker.pending < bestWorker.pending) {
        bestWorker = worker;
      }
    }
    return bestWorker;
  }

  private async getContexts(
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<LlamaEmbeddingContext[]>> {
    if (this.workers.some((worker) => worker.context.disposed)) {
      await this.dispose();
    }
    if (this.workers.length > 0) {
      return Promise.resolve({
        ok: true,
        value: this.workers.map((worker) => worker.context),
      });
    }

    if (this.contextsPromise) {
      return this.contextsPromise;
    }

    this.contextsPromise = this.createContexts(options);
    return this.contextsPromise;
  }

  private resolveTargetPoolSize(llama: Llama): number {
    return resolveEmbeddingContextPoolSize({
      gpu: llama.gpu,
      cpuMathCores: llama.cpuMathCores,
    });
  }

  private threadsPerContext = 0;

  private resolveThreadsPerContext(llama: Llama, poolSize: number): number {
    if (llama.gpu !== false) {
      return 0;
    }

    const override = resolveThreadsPerContextOverride();
    if (override !== undefined) {
      return override;
    }

    return Math.max(1, Math.floor(Math.max(1, llama.cpuMathCores) / poolSize));
  }

  private async createContexts(
    options?: NativeEvaluationOptions
  ): Promise<LlmResult<LlamaEmbeddingContext[]>> {
    const lifecycleVersion = this.lifecycleVersion;
    const model = await this.manager.loadModel(
      this.modelPath,
      this.modelUri,
      "embed",
      options?.signal
    );
    if (!model.ok) {
      this.contextsPromise = null;
      return model;
    }

    try {
      const llamaModel = model.value.model as LlamaModel;
      const llama = await this.manager.getLlama();
      this.embeddingContextSize =
        resolveEmbeddingContextSizeOverride() ?? DEFAULT_EMBEDDING_CONTEXT_SIZE;
      const targetPoolSize = this.resolveTargetPoolSize(llama);
      const threadsPerContext = this.resolveThreadsPerContext(
        llama,
        targetPoolSize
      );
      this.threadsPerContext = threadsPerContext;
      const contextOptions =
        llama.gpu === false
          ? {
              contextSize: this.embeddingContextSize,
              threads: threadsPerContext,
            }
          : { contextSize: this.embeddingContextSize };
      const contexts: LlamaEmbeddingContext[] = [];

      for (let i = 0; i < targetPoolSize; i += 1) {
        try {
          const context = await llamaModel.createEmbeddingContext({
            ...contextOptions,
            createSignal: options?.signal,
          });
          contexts.push(context);
        } catch (error) {
          if (options?.signal?.aborted) {
            await Promise.allSettled(
              contexts.map((context) => context.dispose())
            );
            this.contextsPromise = null;
            throw error;
          }
          if (contexts.length === 0) {
            this.contextsPromise = null;
            return {
              ok: false,
              error: inferenceFailedError(this.modelUri, error),
            };
          }
          break;
        }
      }

      if (lifecycleVersion !== this.lifecycleVersion) {
        for (const context of contexts) {
          try {
            await context.dispose();
          } catch {
            // Ignore disposal errors
          }
        }
        return {
          ok: false,
          error: inferenceFailedError(
            this.modelUri,
            new Error("Embedding context disposed during initialization")
          ),
        };
      }

      this.workers = contexts.map((context) => ({ context, pending: 0 }));
      this.llamaModel = llamaModel as TokenizingModel;

      const size = llamaModel.embeddingVectorSize;
      if (this.dims === null && typeof size === "number" && size > 0) {
        this.dims = size;
      }

      return { ok: true, value: contexts };
    } catch (e) {
      this.contextsPromise = null;
      return { ok: false, error: inferenceFailedError(this.modelUri, e) };
    }
  }

  private truncateForEmbedding(
    text: string,
    mode: "single" | "batch"
  ): LlmResult<{ input: EmbeddingInput }> {
    const model = this.llamaModel;
    const modelLimit =
      typeof model?.trainContextSize === "number" &&
      Number.isFinite(model.trainContextSize) &&
      model.trainContextSize > 0
        ? Math.floor(model.trainContextSize)
        : undefined;
    if (!model) {
      return { ok: true, value: { input: text } };
    }

    const rawLimit =
      modelLimit === undefined
        ? this.embeddingContextSize
        : Math.min(modelLimit, this.embeddingContextSize);
    const limit = Math.max(1, rawLimit - 4);
    try {
      const tokens = model.tokenize(text);
      if (tokens.length <= limit) {
        return {
          ok: true,
          value: { input: tokens as EmbeddingInput },
        };
      }

      const truncatedTokens = tokens.slice(0, limit);
      const shouldWarn =
        mode === "single"
          ? !this.warnedSingleTruncation
          : !this.warnedBatchTruncation;
      if (shouldWarn) {
        if (mode === "single") {
          this.warnedSingleTruncation = true;
        } else {
          this.warnedBatchTruncation = true;
        }
        console.warn(
          `[llama] Truncated embedding input from ${tokens.length} to ${limit} tokens`
        );
      }
      return {
        ok: true,
        value: { input: truncatedTokens as EmbeddingInput },
      };
    } catch (error) {
      return { ok: false, error: inferenceFailedError(this.modelUri, error) };
    }
  }
}