import { cpus } from "node:os";

import { commercialApiHold } from "./model-policy.js";

export function isLocalLlmDisabled(_env: NodeJS.ProcessEnv = process.env): boolean {
  return true;
}

export function resolveLlamaGpuMode(_env: NodeJS.ProcessEnv = process.env): "cpu" | "auto" {
  return "cpu";
}

export function isQwen3EmbeddingModel(modelUri: string): boolean {
  return /qwen.*embed/i.test(modelUri) || /embed.*qwen/i.test(modelUri);
}

export function formatQueryForEmbedding(query: string, modelUri?: string): string {
  if (modelUri && isQwen3EmbeddingModel(modelUri)) {
    return `Instruct: Retrieve relevant documents for the given query\nQuery: ${query}`;
  }
  return `task: search result | query: ${query}`;
}

export function formatDocForEmbedding(text: string, title?: string, modelUri?: string): string {
  if (modelUri && isQwen3EmbeddingModel(modelUri)) {
    return title ? `${title}\n${text}` : text;
  }
  return `title: ${title || "none"} | text: ${text}`;
}

export type TokenLogProb = { token: string; logprob: number };
export type EmbeddingResult = { embedding: number[]; model: string };
export type GenerateResult = {
  text: string;
  model: string;
  logprobs?: TokenLogProb[];
  done: boolean;
};
export type RerankDocumentResult = { file: string; score: number; index: number };
export type RerankResult = { results: RerankDocumentResult[]; model: string };
export type ModelInfo = { name: string; exists: boolean; path?: string };
export type EmbedOptions = { model?: string; isQuery?: boolean; title?: string };
export type GenerateOptions = { model?: string; maxTokens?: number; temperature?: number };
export type RerankOptions = { model?: string };
export type LLMSessionOptions = { maxDuration?: number; signal?: AbortSignal; name?: string };
export type QueryType = "lex" | "vec" | "hyde";
export type Queryable = { type: QueryType; text: string };
export type RerankDocument = { file: string; text: string; title?: string };

export interface ILLMSession {
  embed(text: string, options?: EmbedOptions): Promise<EmbeddingResult | null>;
  embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]>;
  expandQuery(query: string, options?: { context?: string; includeLexical?: boolean }): Promise<Queryable[]>;
  rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise<RerankResult>;
  readonly isValid: boolean;
  readonly signal: AbortSignal;
}

export const LFM2_GENERATE_MODEL = "commercial-api-required";
export const LFM2_INSTRUCT_MODEL = "commercial-api-required";
export const DEFAULT_EMBED_MODEL_URI = "commercial-api:embedding-unconfigured";
export const DEFAULT_RERANK_MODEL_URI = "commercial-api:rerank-unconfigured";
export const DEFAULT_GENERATE_MODEL_URI = "commercial-api:generation-unconfigured";
export const DEFAULT_MODEL_CACHE_DIR = "commercial-api:no-local-cache";

export type PullResult = {
  model: string;
  path: string;
  sizeBytes: number;
  refreshed: boolean;
};

export async function pullModels(
  _models: string[],
  _options: { refresh?: boolean; cacheDir?: string } = {},
): Promise<PullResult[]> {
  throw commercialApiHold("local model downloads are disabled");
}

export interface LLM {
  embed(text: string, options?: EmbedOptions): Promise<EmbeddingResult | null>;
  embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]>;
  generate(prompt: string, options?: GenerateOptions): Promise<GenerateResult | null>;
  modelExists(modelUri: string): Promise<ModelInfo>;
  expandQuery(query: string, options?: { context?: string; includeLexical?: boolean; intent?: string }): Promise<Queryable[]>;
  rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise<RerankResult>;
  dispose(): Promise<void>;
}

export type LlamaCppConfig = {
  embedModel?: string;
  generateModel?: string;
  rerankModel?: string;
  modelCacheDir?: string;
  expandContextSize?: number;
  inactivityTimeoutMs?: number;
  disposeModelsOnInactivity?: boolean;
};

function learnedModelHold(operation: string): never {
  throw commercialApiHold(`${operation} requires an approved commercial API adapter`);
}

/** Historical SDK name retained as a fail-closed compatibility adapter. */
export class LlamaCpp implements LLM {
  static readonly EMBED_CONTEXT_SIZE = 2048;
  static readonly RERANK_CONTEXT_SIZE = 2048;
  static readonly RERANK_TARGET_DOCS_PER_CONTEXT = 32;
  static readonly RERANK_TEMPLATE_OVERHEAD = 32;

  readonly embedModelName: string;

  constructor(config: LlamaCppConfig = {}) {
    this.embedModelName = config.embedModel ?? process.env.QMD_EMBED_MODEL_ID ?? DEFAULT_EMBED_MODEL_URI;
  }

  async tokenize(_text: string): Promise<readonly number[]> {
    return learnedModelHold("tokenization");
  }

  async countTokens(_text: string): Promise<number> {
    return learnedModelHold("token counting");
  }

  async detokenize(_tokens: readonly number[]): Promise<string> {
    return learnedModelHold("detokenization");
  }

  async embed(_text: string, _options: EmbedOptions = {}): Promise<EmbeddingResult | null> {
    return learnedModelHold("embedding");
  }

  async embedBatch(_texts: string[], _options: EmbedOptions = {}): Promise<(EmbeddingResult | null)[]> {
    return learnedModelHold("batch embedding");
  }

  async generate(_prompt: string, _options: GenerateOptions = {}): Promise<GenerateResult | null> {
    return learnedModelHold("generation");
  }

  async modelExists(modelUri: string): Promise<ModelInfo> {
    return { name: modelUri, exists: false };
  }

  async expandQuery(
    _query: string,
    _options: { context?: string; includeLexical?: boolean; intent?: string } = {},
  ): Promise<Queryable[]> {
    return learnedModelHold("query expansion");
  }

  async rerank(
    _query: string,
    _documents: RerankDocument[],
    _options: RerankOptions = {},
  ): Promise<RerankResult> {
    return learnedModelHold("reranking");
  }

  async getDeviceInfo(): Promise<{
    gpu: string | false;
    gpuOffloading: boolean;
    gpuDevices: string[];
    vram?: { total: number; used: number; free: number };
    cpuCores: number;
  }> {
    return { gpu: false, gpuOffloading: false, gpuDevices: [], cpuCores: cpus().length };
  }

  async unloadIdleResources(): Promise<void> {}
  async dispose(): Promise<void> {}
}

export class SessionReleasedError extends Error {
  constructor(message = "LLM session has been released or aborted") {
    super(message);
    this.name = "SessionReleasedError";
  }
}

class CommercialHoldSession implements ILLMSession {
  private released = false;
  private readonly controller = new AbortController();

  constructor(private readonly llm: LlamaCpp, options: LLMSessionOptions = {}) {
    if (options.signal?.aborted) this.controller.abort(options.signal.reason);
    options.signal?.addEventListener("abort", () => this.controller.abort(options.signal?.reason), { once: true });
  }

  get isValid(): boolean {
    return !this.released && !this.controller.signal.aborted;
  }

  get signal(): AbortSignal {
    return this.controller.signal;
  }

  release(): void {
    this.released = true;
    this.controller.abort(new SessionReleasedError());
  }

  private assertValid(): void {
    if (!this.isValid) throw new SessionReleasedError();
  }

  async embed(text: string, options?: EmbedOptions): Promise<EmbeddingResult | null> {
    this.assertValid();
    return this.llm.embed(text, options);
  }

  async embedBatch(texts: string[], options?: EmbedOptions): Promise<(EmbeddingResult | null)[]> {
    this.assertValid();
    return this.llm.embedBatch(texts, options);
  }

  async expandQuery(query: string, options?: { context?: string; includeLexical?: boolean }): Promise<Queryable[]> {
    this.assertValid();
    return this.llm.expandQuery(query, options);
  }

  async rerank(query: string, documents: RerankDocument[], options?: RerankOptions): Promise<RerankResult> {
    this.assertValid();
    return this.llm.rerank(query, documents, options);
  }
}

let defaultLlamaCpp: LlamaCpp | null = null;

export function getDefaultLlamaCpp(): LlamaCpp {
  defaultLlamaCpp ??= new LlamaCpp();
  return defaultLlamaCpp;
}

export function setDefaultLlamaCpp(llm: LlamaCpp | null): void {
  defaultLlamaCpp = llm;
}

export async function disposeDefaultLlamaCpp(): Promise<void> {
  if (defaultLlamaCpp) await defaultLlamaCpp.dispose();
  defaultLlamaCpp = null;
}

export async function withLLMSession<T>(
  fn: (session: ILLMSession) => Promise<T>,
  options: LLMSessionOptions = {},
): Promise<T> {
  return withLLMSessionForLlm(getDefaultLlamaCpp(), fn, options);
}

export async function withLLMSessionForLlm<T>(
  llm: LlamaCpp,
  fn: (session: ILLMSession) => Promise<T>,
  options: LLMSessionOptions = {},
): Promise<T> {
  const session = new CommercialHoldSession(llm, options);
  try {
    return await fn(session);
  } finally {
    session.release();
  }
}

export function canUnloadLLM(): boolean {
  return true;
}
