/**
 * SDK embedding helpers.
 *
 * @module src/sdk/embed
 */

import type { Database } from "bun:sqlite";

import type { Config } from "../config/types";
import type { LlmAdapter } from "../llm/nodeLlamaCpp/adapter";
import type { EmbeddingPort } from "../llm/types";
import type { SqliteAdapter } from "../store/sqlite/adapter";
import type { StoreResult } from "../store/types";
import type {
  BacklogItem,
  VectorIndexPort,
  VectorStatsPort,
} from "../store/vector";
import type { GnoEmbedOptions, GnoEmbedResult } from "./types";

import { embedBacklog, prepareEmbeddingBacklog } from "../embed/backlog";
import { getEmbeddingFingerprint } from "../embed/fingerprint";
import {
  chunkRetryKey,
  embedAndStoreBatch,
  MAX_EMBED_CHUNK_ATTEMPTS,
} from "../embed/retry";
import { countVariantBacklog } from "../embed/variant-plan";
import { resolveModelUri } from "../llm/registry";
import { err, ok } from "../store/types";
import { createVectorIndexPort, createVectorStatsPort } from "../store/vector";
import { sdkError } from "./errors";

interface EmbedRuntimeOptions {
  config: Config;
  store: SqliteAdapter;
  llm: LlmAdapter;
  downloadPolicy?: import("../llm/policy").DownloadPolicy;
}

function getActiveChunkCount(db: Database): Promise<StoreResult<number>> {
  try {
    const result = db
      .prepare(
        `
        SELECT COUNT(*) as count FROM content_chunks c
        WHERE EXISTS (
          SELECT 1 FROM documents d
          WHERE d.mirror_hash = c.mirror_hash AND d.active = 1
        )
      `
      )
      .get() as { count: number };
    return Promise.resolve(ok(result.count));
  } catch (cause) {
    return Promise.resolve(
      err(
        "QUERY_FAILED",
        cause instanceof Error ? cause.message : "Failed to count chunks",
        cause
      )
    );
  }
}

function getActiveChunks(
  db: Database,
  limit: number,
  after?: { mirrorHash: string; seq: number }
): Promise<StoreResult<BacklogItem[]>> {
  try {
    const sql = after
      ? `
        SELECT c.mirror_hash as mirrorHash, c.seq, c.text,
          (SELECT d.title FROM documents d WHERE d.mirror_hash = c.mirror_hash AND d.active = 1 ORDER BY d.id LIMIT 1) as title,
          'force' as reason
        FROM content_chunks c
        WHERE EXISTS (
          SELECT 1 FROM documents d
          WHERE d.mirror_hash = c.mirror_hash AND d.active = 1
        )
        AND (c.mirror_hash > ? OR (c.mirror_hash = ? AND c.seq > ?))
        ORDER BY c.mirror_hash, c.seq
        LIMIT ?
      `
      : `
        SELECT c.mirror_hash as mirrorHash, c.seq, c.text,
          (SELECT d.title FROM documents d WHERE d.mirror_hash = c.mirror_hash AND d.active = 1 ORDER BY d.id LIMIT 1) as title,
          'force' as reason
        FROM content_chunks c
        WHERE EXISTS (
          SELECT 1 FROM documents d
          WHERE d.mirror_hash = c.mirror_hash AND d.active = 1
        )
        ORDER BY c.mirror_hash, c.seq
        LIMIT ?
      `;

    const params = after
      ? [after.mirrorHash, after.mirrorHash, after.seq, limit]
      : [limit];
    const result = db.prepare(sql).all(...params) as BacklogItem[];
    return Promise.resolve(ok(result));
  } catch (cause) {
    return Promise.resolve(
      err(
        "QUERY_FAILED",
        cause instanceof Error ? cause.message : "Failed to get chunks",
        cause
      )
    );
  }
}

async function forceEmbedAll(
  db: Database,
  embedPort: EmbeddingPort,
  vectorIndex: VectorIndexPort,
  modelUri: string,
  batchSize: number
): Promise<{ embedded: number; errors: number; contentionErrors: number }> {
  let embedded = 0;
  let errors = 0;
  let contentionErrors = 0;
  let cursor: { mirrorHash: string; seq: number } | undefined;
  const retryQueue = new Map<string, { item: BacklogItem; attempts: number }>();
  const embedFingerprint = getEmbeddingFingerprint({
    modelUri,
    dimensions: vectorIndex.dimensions,
  });

  const enqueueRetryItems = (items: BacklogItem[], attempts: number): void => {
    for (const item of items) {
      const key = chunkRetryKey(item);
      const existing = retryQueue.get(key);
      retryQueue.set(key, {
        item,
        attempts: Math.max(existing?.attempts ?? 0, attempts),
      });
    }
  };

  const drainRetryQueue = async (): Promise<number> => {
    if (retryQueue.size === 0) {
      return 0;
    }

    let retryEmbedded = 0;
    const entries = [...retryQueue.values()].filter(
      (entry) => entry.attempts < MAX_EMBED_CHUNK_ATTEMPTS
    );

    for (let idx = 0; idx < entries.length; idx += batchSize) {
      const slice = entries.slice(idx, idx + batchSize);
      for (const entry of slice) {
        retryQueue.delete(chunkRetryKey(entry.item));
        entry.attempts += 1;
      }

      const retryResult = await embedAndStoreBatch({
        embedPort,
        vectorIndex,
        items: slice.map((entry) => entry.item),
        modelUri,
        embedFingerprint,
      });
      embedded += retryResult.embedded;
      errors += retryResult.errors;
      contentionErrors += retryResult.contentionErrors;
      retryEmbedded += retryResult.embedded;

      const retryByKey = new Set(
        retryResult.retryItems.map((item) => chunkRetryKey(item))
      );
      for (const entry of slice) {
        if (!retryByKey.has(chunkRetryKey(entry.item))) {
          continue;
        }
        if (entry.attempts >= MAX_EMBED_CHUNK_ATTEMPTS) {
          errors += 1;
        } else {
          retryQueue.set(chunkRetryKey(entry.item), entry);
        }
      }
    }

    return retryEmbedded;
  };

  while (true) {
    const batchResult = await getActiveChunks(db, batchSize, cursor);
    if (!batchResult.ok) {
      throw sdkError("STORE", batchResult.error.message, {
        cause: batchResult.error.cause,
      });
    }

    const batch = batchResult.value;
    if (batch.length === 0) {
      break;
    }

    const lastItem = batch.at(-1);
    if (lastItem) {
      cursor = { mirrorHash: lastItem.mirrorHash, seq: lastItem.seq };
    }

    const beforeEmbedded = embedded;
    const embedResult = await embedAndStoreBatch({
      embedPort,
      vectorIndex,
      items: batch,
      modelUri,
      embedFingerprint,
    });
    embedded += embedResult.embedded;
    errors += embedResult.errors;
    contentionErrors += embedResult.contentionErrors;
    enqueueRetryItems(embedResult.retryItems, 1);

    if (embedded > beforeEmbedded) {
      await drainRetryQueue();
    }
  }

  await drainRetryQueue();
  if (retryQueue.size > 0) {
    errors += retryQueue.size;
    retryQueue.clear();
  }

  if (vectorIndex.vecDirty) {
    const syncResult = await vectorIndex.syncVecIndex();
    if (syncResult.ok) {
      vectorIndex.vecDirty = false;
    }
  }

  return { embedded, errors, contentionErrors };
}

export async function runEmbed(
  runtime: EmbedRuntimeOptions,
  options: GnoEmbedOptions = {}
): Promise<GnoEmbedResult> {
  const batchSize = options.batchSize ?? 32;
  const force = options.force ?? false;
  const dryRun = options.dryRun ?? false;
  const modelUri = resolveModelUri(
    runtime.config,
    "embed",
    options.model,
    options.collection
  );
  const db = runtime.store.getRawDb();
  const stats: VectorStatsPort = createVectorStatsPort(db);

  let totalToEmbed = 0;
  const embedResult = await runtime.llm.createEmbeddingPort(modelUri, {
    egressCollections: options.collection ? [options.collection] : "all",
    policy: runtime.downloadPolicy,
  });
  if (!embedResult.ok) {
    throw sdkError("MODEL", embedResult.error.message, {
      cause: embedResult.error.cause,
    });
  }

  const embedPort = embedResult.value;
  try {
    const initializedPort = await embedPort.init();
    if (!initializedPort.ok)
      throw sdkError("MODEL", initializedPort.error.message);
    let dimensions = embedPort.dimensions();
    if (!embedPort.getIdentity?.()) {
      const probeResult = await embedPort.embed("dimension probe");
      if (!probeResult.ok)
        throw sdkError("MODEL", probeResult.error.message, {
          cause: probeResult.error.cause,
        });
      dimensions = probeResult.value.length;
    }

    const vectorResult = await createVectorIndexPort(db, {
      model: modelUri,
      dimensions,
    });
    if (!vectorResult.ok) {
      throw sdkError("STORE", vectorResult.error.message, {
        cause: vectorResult.error.cause,
      });
    }

    const vectorIndex = vectorResult.value;
    const prepared = await prepareEmbeddingBacklog({
      statsPort: stats,
      embedPort,
      vectorIndex,
      modelUri,
      collection: options.collection,
      batchSize,
      force,
    });
    if (!prepared.ok) throw sdkError("STORE", prepared.error.message);
    if (prepared.value.variantStore) {
      const count = countVariantBacklog(prepared.value);
      const startedAt = Date.now();
      if (dryRun)
        return {
          embedded: count,
          errors: 0,
          duration: 0,
          model: modelUri,
          searchAvailable: prepared.value.variantStore.searchAvailable,
        };
      const processed = await embedBacklog(prepared.value);
      if (!processed.ok) throw sdkError("STORE", processed.error.message);
      if (processed.value.syncError)
        throw sdkError("STORE", processed.value.syncError);
      return {
        embedded: processed.value.embedded,
        errors: processed.value.errors,
        contentionErrors: processed.value.contentionErrors ?? 0,
        duration: (Date.now() - startedAt) / 1000,
        model: modelUri,
        searchAvailable: prepared.value.variantStore.searchAvailable,
      };
    }
    if (force) {
      const count = await getActiveChunkCount(db);
      if (!count.ok) throw sdkError("STORE", count.error.message);
      totalToEmbed = count.value;
      if (dryRun || !totalToEmbed)
        return {
          embedded: totalToEmbed,
          errors: 0,
          duration: 0,
          model: modelUri,
          searchAvailable: vectorIndex.searchAvailable,
        };
    }
    if (!force) {
      const embedFingerprint = getEmbeddingFingerprint({
        modelUri,
        dimensions: vectorIndex.dimensions,
      });
      const backlogResult = await stats.countBacklog(
        modelUri,
        embedFingerprint,
        {
          collection: options.collection,
        }
      );
      if (!backlogResult.ok) {
        throw sdkError("STORE", backlogResult.error.message, {
          cause: backlogResult.error.cause,
        });
      }

      totalToEmbed = backlogResult.value;
      if (totalToEmbed === 0 || dryRun) {
        return {
          embedded: totalToEmbed,
          errors: 0,
          duration: 0,
          model: modelUri,
          searchAvailable: vectorIndex.searchAvailable,
        };
      }
    }

    const startedAt = Date.now();
    let result: { embedded: number; errors: number; contentionErrors: number };
    if (force) {
      result = await forceEmbedAll(
        db,
        embedPort,
        vectorIndex,
        modelUri,
        batchSize
      );
    } else {
      const processed = await embedBacklog({
        statsPort: stats,
        embedPort,
        vectorIndex,
        collection: options.collection,
        modelUri,
        batchSize,
      });
      if (!processed.ok) {
        throw sdkError("STORE", processed.error.message, {
          cause: processed.error.cause,
        });
      }
      result = {
        embedded: processed.value.embedded,
        errors: processed.value.errors,
        contentionErrors: processed.value.contentionErrors ?? 0,
      };
    }

    return {
      embedded: result.embedded,
      errors: result.errors,
      contentionErrors: result.contentionErrors,
      duration: (Date.now() - startedAt) / 1000,
      model: modelUri,
      searchAvailable: vectorIndex.searchAvailable,
    };
  } finally {
    await embedPort.dispose();
  }
}