import { env } from "cloudflare:workers";
import { batch, chunks, first, placeholders, run, stmt } from "../db";
import { loadRepoById } from "../repos";
import { branchSha, listRefs } from "../git/refs";
import { withStore } from "../git/store";
import type { GitStore } from "../git/store";
import { diffTreesBounded } from "../git/diff";
import type { RepoRow } from "~/lib/types";

const FILES_PER_MESSAGE = 100;
const MAX_INDEXED_FILES = 5000;
const MAX_INDEXED_BYTES = 128 * 1024;
const BINARY_SNIFF = 8000;
const READ_CONCURRENCY = 8;
const SKIPPED_MODES = new Set(["120000", "160000"]);
const decoder = new TextDecoder("utf-8", { fatal: true });

type IndexJob = { commit: string; tree: string; base: string | null; offset: number };

export type CodeIndexMessage = { kind: "code.index"; repoId: string; job?: IndexJob };

type Plan = { full: boolean; upserts: Array<{ path: string; sha: string }>; deletes: string[] };

async function planChanges(store: GitStore, base: string | null, tree: string): Promise<Plan> {
  if (base) {
    const { changes, truncated } = await diffTreesBounded(store, base, tree);
    if (!truncated) {
      const upserts: Plan["upserts"] = [];
      const deletes: string[] = [];
      for (const change of changes) {
        if (change.oldPath) deletes.push(change.oldPath);
        if (change.status === "deleted" || !change.newSha || SKIPPED_MODES.has(change.newMode ?? "")) {
          deletes.push(change.path);
        } else {
          upserts.push({ path: change.path, sha: change.newSha });
        }
      }
      return { full: false, upserts, deletes };
    }
  }
  const { entries } = await store.walk(tree, MAX_INDEXED_FILES, { withContents: true });
  return {
    full: true,
    upserts: entries
      .filter((entry) => !SKIPPED_MODES.has(entry.mode))
      .map((entry) => ({ path: entry.path, sha: entry.sha })),
    deletes: [],
  };
}

async function readText(store: GitStore, sha: string): Promise<string | null> {
  const bytes = await store.blob(sha);
  if (!bytes || bytes.byteLength > MAX_INDEXED_BYTES || bytes.subarray(0, BINARY_SNIFF).includes(0)) return null;
  try {
    return decoder.decode(bytes);
  } catch {
    return null;
  }
}

function upsertStatements(repoId: string, path: string, sha: string, text: string): D1PreparedStatement[] {
  return [
    stmt(
      `INSERT INTO code_files (repo_id, path, blob_sha) SELECT ?1, ?2, ?3
       WHERE (SELECT COUNT(*) FROM code_files WHERE repo_id = ?1) < ?4
          OR EXISTS (SELECT 1 FROM code_files WHERE repo_id = ?1 AND path = ?2)
       ON CONFLICT (repo_id, path) DO UPDATE SET blob_sha = excluded.blob_sha`,
      repoId,
      path,
      sha,
      MAX_INDEXED_FILES,
    ),
    stmt(
      "DELETE FROM code_fts WHERE rowid = (SELECT id FROM code_files WHERE repo_id = ? AND path = ? AND blob_sha = ?)",
      repoId,
      path,
      sha,
    ),
    stmt(
      `INSERT INTO code_fts (rowid, path, content)
       SELECT id, path, ? FROM code_files WHERE repo_id = ? AND path = ? AND blob_sha = ?`,
      text,
      repoId,
      path,
      sha,
    ),
  ];
}

async function applyBatch(store: GitStore, repoId: string, items: Plan["upserts"]): Promise<void> {
  const statements: D1PreparedStatement[] = [];
  for (let index = 0; index < items.length; index += READ_CONCURRENCY) {
    const group = items.slice(index, index + READ_CONCURRENCY);
    const texts = await Promise.all(group.map((item) => readText(store, item.sha)));
    group.forEach((item, position) => {
      const text = texts[position];
      if (text === null) {
        statements.push(stmt("DELETE FROM code_files WHERE repo_id = ? AND path = ?", repoId, item.path));
      } else {
        statements.push(...upsertStatements(repoId, item.path, item.sha, text));
      }
    });
  }
  await batch(statements);
}

async function applyDeletes(repoId: string, plan: Plan): Promise<void> {
  if (plan.full) await run("DELETE FROM code_files WHERE repo_id = ?", repoId);
  for (const part of chunks([...new Set(plan.deletes)])) {
    await run(`DELETE FROM code_files WHERE repo_id = ? AND path IN (${placeholders(part.length)})`, repoId, ...part);
  }
}

async function finish(repoId: string, job: IndexJob, plan: Plan): Promise<void> {
  const count = await first<{ total: number }>("SELECT COUNT(*) AS total FROM code_files WHERE repo_id = ?", repoId);
  const total = count?.total ?? 0;
  await run(
    `INSERT INTO code_index_state (repo_id, commit_sha, tree_sha, file_count, truncated, updated_at)
     VALUES (?, ?, ?, ?, ?, ?)
     ON CONFLICT (repo_id) DO UPDATE SET commit_sha = excluded.commit_sha, tree_sha = excluded.tree_sha,
       file_count = excluded.file_count, truncated = excluded.truncated, updated_at = excluded.updated_at`,
    repoId,
    job.commit,
    job.tree,
    total,
    total >= MAX_INDEXED_FILES || (plan.full && plan.upserts.length >= MAX_INDEXED_FILES) ? 1 : 0,
    Date.now(),
  );
}

/**
 * Indexes the default branch head into `code_fts`. Each message handles at most FILES_PER_MESSAGE
 * blobs and re-enqueues itself with an offset, so a large push becomes a bounded chain of messages.
 */
export async function indexRepoCode(message: CodeIndexMessage): Promise<void> {
  const repo = await loadRepoById(message.repoId);
  if (!repo || repo.status !== "ready") return;
  const refs = await listRefs(repo);
  const head = branchSha(refs, repo.default_branch);
  if (!head) return;
  if (message.job && message.job.commit !== head) return;
  await withStore(repo.artifacts_name, async (store) => {
    let job = message.job;
    if (!job) {
      const state = await first<{ commit_sha: string; tree_sha: string }>(
        "SELECT commit_sha, tree_sha FROM code_index_state WHERE repo_id = ?",
        repo.id,
      );
      if (state?.commit_sha === head) return;
      const commit = await store.commit(head);
      if (!commit) return;
      job = { commit: head, tree: commit.treeSha, base: state?.tree_sha ?? null, offset: 0 };
    }
    const plan = await planChanges(store, job.base, job.tree);
    if (job.offset === 0) await applyDeletes(repo.id, plan);
    const next = job.offset + FILES_PER_MESSAGE;
    await applyBatch(store, repo.id, plan.upserts.slice(job.offset, next));
    if (next < plan.upserts.length) {
      const continuation: CodeIndexMessage = { kind: "code.index", repoId: repo.id, job: { ...job, offset: next } };
      await env.EVENTS.send(continuation);
      return;
    }
    await finish(repo.id, job, plan);
  });
}

export async function requestCodeIndex(repoId: string): Promise<void> {
  const message: CodeIndexMessage = { kind: "code.index", repoId };
  await env.EVENTS.send(message);
}

const REQUEST_DEBOUNCE_SECONDS = 300;

/** Queues an index build when the default branch head is not the indexed commit; debounced so repeat searches don't stack jobs. */
export async function ensureCodeIndex(repo: RepoRow, indexedSha: string | null): Promise<boolean> {
  if (repo.status !== "ready") return false;
  const head = branchSha(await listRefs(repo), repo.default_branch);
  if (!head || head === indexedSha) return false;
  const key = `code-index-requested:${repo.id}:${head}`;
  if (await env.KV.get(key)) return true;
  await env.KV.put(key, "1", { expirationTtl: REQUEST_DEBOUNCE_SECONDS });
  await requestCodeIndex(repo.id);
  return true;
}