import fs from "node:fs";
import path from "node:path";
import type { EmbeddingModelV2, EmbeddingModelV2Embedding } from "@ai-sdk/provider";

export const DEFAULT_LOCAL_EMBEDDING_MODEL = "Xenova/bge-large-en-v1.5";
const DEFAULT_QUERY_PREFIX = "Represent this code for search:";
const DEFAULT_QUERY_MAX_CHARS = 512;
const DEFAULT_MAX_BATCH = 64;

// Define a type that matches the actual pipeline return value
type FeatureExtractionPipeline = (
  texts: string[],
  options?: { pooling?: string; normalize?: boolean }
) => Promise<unknown>;

export interface LocalEmbeddingOptions {
  modelId?: string;
  modelDir?: string;
  cacheDir?: string;
  allowRemoteModels?: boolean;
  queryPrefix?: string;
  queryPrefixMaxChars?: number;
}

let cachedPipelineKey: string | null = null;
let cachedPipelinePromise: Promise<FeatureExtractionPipeline> | null = null;

function resolveCacheDir(override?: string): string {
  if (override) return override;
  if (process.env.EMBEDDING_CACHE_DIR) return process.env.EMBEDDING_CACHE_DIR;
  const basePath = process.env.LOCAL_DATA_PATH || path.join(process.cwd(), ".local-data");
  return path.join(basePath, "transformers-cache");
}

function resolveModelDir(override?: string): string | undefined {
  if (override) return override;
  if (process.env.EMBEDDING_MODEL_DIR) return process.env.EMBEDDING_MODEL_DIR;

  // Fallback: read from settings.json (for Next.js dev server which doesn't have Electron env vars)
  try {
    const basePath = process.env.LOCAL_DATA_PATH || path.join(process.cwd(), ".local-data");
    const settingsPath = path.join(basePath, "settings.json");
    if (fs.existsSync(settingsPath)) {
      const settings = JSON.parse(fs.readFileSync(settingsPath, "utf-8"));
      if (settings.embeddingModelDir && fs.existsSync(settings.embeddingModelDir)) {
        return settings.embeddingModelDir;
      }
    }
  } catch {
    // Ignore errors reading settings
  }

  return undefined;
}

function resolveQueryPrefix(override?: string | null): string {
  if (override !== undefined) return override ?? "";
  if (process.env.LOCAL_EMBEDDING_QUERY_PREFIX !== undefined) {
    return process.env.LOCAL_EMBEDDING_QUERY_PREFIX;
  }
  return DEFAULT_QUERY_PREFIX;
}

function resolveQueryMaxChars(override?: number): number {
  if (Number.isFinite(override)) return Number(override);
  const envValue = process.env.LOCAL_EMBEDDING_QUERY_MAX_CHARS;
  if (envValue) {
    const parsed = Number(envValue);
    if (Number.isFinite(parsed)) return parsed;
  }
  return DEFAULT_QUERY_MAX_CHARS;
}

function applyQueryPrefix(text: string, prefix: string, maxChars: number): string {
  if (!prefix) return text;
  if (text.length > maxChars) return text;
  return `${prefix} ${text}`.trim();
}

function resolveTokenizerPath(cacheDir: string, modelId: string): string {
  const parts = modelId.split("/");
  return path.join(cacheDir, ...parts, "tokenizer.json");
}

function patchTokenizerMerges(tokenizerPath: string): boolean {
  if (!tokenizerPath || !fs.existsSync(tokenizerPath)) {
    return false;
  }
  const raw = fs.readFileSync(tokenizerPath, "utf8");
  const data = JSON.parse(raw);
  const merges = data?.merges ?? data?.model?.merges;
  if (!Array.isArray(merges) || merges.length === 0) {
    return false;
  }
  if (!Array.isArray(merges[0])) {
    return false;
  }
  const converted = merges.map((pair) => pair.join(" "));
  if (data.merges) {
    data.merges = converted;
  } else if (data.model) {
    data.model.merges = converted;
  }
  fs.writeFileSync(tokenizerPath, JSON.stringify(data));
  return true;
}

async function loadPipelineWithPatch(
  modelId: string,
  cacheDir: string
): Promise<FeatureExtractionPipeline> {
  const { pipeline } = await import("@xenova/transformers");
  try {
    return (await pipeline("feature-extraction", modelId)) as unknown as FeatureExtractionPipeline;
  } catch (error) {
    const message = String(error);
    if (message.includes("x.split is not a function")) {
      const tokenizerPath = resolveTokenizerPath(cacheDir, modelId);
      if (patchTokenizerMerges(tokenizerPath)) {
        return (await pipeline("feature-extraction", modelId)) as unknown as FeatureExtractionPipeline;
      }
    }
    throw error;
  }
}

async function getPipeline(options: LocalEmbeddingOptions): Promise<FeatureExtractionPipeline> {
  const modelId = options.modelId ?? DEFAULT_LOCAL_EMBEDDING_MODEL;
  const cacheDir = resolveCacheDir(options.cacheDir);
  const modelDir = resolveModelDir(options.modelDir);
  const allowRemoteModels = options.allowRemoteModels ?? !modelDir;
  const cacheKey = [modelId, cacheDir, modelDir ?? "", allowRemoteModels ? "remote" : "local"].join("|");

  if (cachedPipelinePromise && cachedPipelineKey === cacheKey) {
    return cachedPipelinePromise;
  }

  cachedPipelineKey = cacheKey;
  cachedPipelinePromise = (async () => {
    const { env } = await import("@xenova/transformers");
    env.cacheDir = cacheDir;
    env.useBrowserCache = false;
    env.allowLocalModels = true;
    env.allowRemoteModels = allowRemoteModels;
    if (modelDir) {
      env.localModelPath = modelDir;
      if (!fs.existsSync(modelDir)) {
        console.warn(`[LocalEmbeddings] Model dir not found: ${modelDir}`);
      }
    }

    return loadPipelineWithPatch(modelId, cacheDir);
  })();

  return cachedPipelinePromise;
}

function toEmbeddings(output: unknown, expectedCount: number): EmbeddingModelV2Embedding[] {
  if (!output) return [];

  if (Array.isArray(output)) {
    return (output as Array<Array<number>>).map((row) =>
      row.map((value) => Number(value))
    );
  }

  const tensor = output as { data?: Float32Array | Float64Array | number[]; dims?: number[] };
  if (!tensor.data || !tensor.dims) {
    return [];
  }

  if (tensor.dims.length !== 2) {
    throw new Error(`Unexpected embedding dims: ${tensor.dims.join("x")}`);
  }

  const [rows, cols] = tensor.dims;
  if (rows !== expectedCount) {
    throw new Error(`Embedding count mismatch: expected ${expectedCount}, got ${rows}`);
  }

  const data = Array.from(tensor.data);
  const embeddings: EmbeddingModelV2Embedding[] = [];
  for (let i = 0; i < rows; i += 1) {
    const start = i * cols;
    embeddings.push(data.slice(start, start + cols));
  }
  return embeddings;
}

export function createLocalEmbeddingModel(
  options: LocalEmbeddingOptions = {}
): EmbeddingModelV2<string> {
  const modelId = options.modelId ?? DEFAULT_LOCAL_EMBEDDING_MODEL;
  const queryPrefix = resolveQueryPrefix(options.queryPrefix ?? null);
  const queryMaxChars = resolveQueryMaxChars(options.queryPrefixMaxChars);

  return {
    specificationVersion: "v2",
    provider: "local",
    modelId,
    maxEmbeddingsPerCall: DEFAULT_MAX_BATCH,
    supportsParallelCalls: false,
    async doEmbed({ values }) {
      if (!values.length) {
        return { embeddings: [] };
      }

      const texts = values.map((value) => applyQueryPrefix(String(value), queryPrefix, queryMaxChars));
      const pipeline = await getPipeline({ ...options, modelId });
      const output = await pipeline(texts, { pooling: "mean", normalize: true });
      const embeddings = toEmbeddings(output, values.length);

      return {
        embeddings,
        usage: { tokens: 0, totalTokens: 0 },
        warnings: [],
      };
    },
  };
}
