import { DurableObject, env } from "cloudflare:workers";
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 { run } from "../db";
import { MAX_POINTER_BYTES, parseLfsPointer } from "../lfs/pointer";
import { parseAttributes } from "~/lib/gitattributes";
import { languageOf, summarizeLanguages } from "~/lib/languages";
import { CODE_PAGE_SIZE, matchingLines } from "~/lib/code-search";
import type { CodeSearchLine } from "~/lib/code-intel";
import { compileQuery, matchLines, requiredLiteral } from "~/lib/text-search";
import type { LineMatch, TextQuery } from "~/lib/text-search";

const FILES_PER_PASS = 100;
const MAX_INDEXED_FILES = 5000;
const MAX_INDEXED_BYTES = 128 * 1024;
/** Files whose sizes feed the language breakdown (any text file, indexed or not). */
const MAX_SIZED_FILES = 20_000;
const MAX_ATTRIBUTES_BYTES = 64 * 1024;
const MAX_ATTRIBUTE_RULES = 500;
/** Bumped when the index gains data (2: file sizes for languages); an older index is rebuilt once. */
const INDEX_SCHEMA = 2;
const BINARY_SNIFF = 8000;
const READ_CONCURRENCY = 8;
const RETRY_ON_ERROR_MS = 60_000;
/** Counting stops here so a broad query stays cheap. */
const COUNT_CAP = 1000;
const LINES_PER_HIT = 3;
const SNIPPET_TOKENS = 24;
const SNIPPET = `snippet(code_fts, 1, char(1), char(2), '…', ${SNIPPET_TOKENS})`;
const SKIPPED_MODES = new Set(["120000", "160000"]);
const decoder = new TextDecoder("utf-8", { fatal: true });

const STATE_KEY = "state";
/** The language summary last written to the repository row, so unchanged ones skip the write. */
const LANGUAGES_KEY = "languages";
const JOB_KEY = "job";

/** `schema` is the index schema the job was planned under; a job left over from an older one finishes as stale. */
type Job = { commit: string; tree: string; base: string | null; offset: number; schema?: number };
type IndexedState = {
  commitSha: string;
  treeSha: string;
  fileCount: number;
  truncated: boolean;
  updatedAt: number;
  schema?: number;
};
type Plan = { full: boolean; upserts: Array<{ path: string; sha: string }>; deletes: string[] };

export type RepoSearchStatus = { commitSha: string; fileCount: number; truncated: boolean } | null;

export type RepoSearchHit = {
  path: string;
  snippet: string;
  lines: CodeSearchLine[];
  matchCount: number;
  /** FTS5 bm25 rank: lower is better, comparable enough to merge results across repositories. */
  rank: number;
};

export type IndexedGrep = {
  files: Array<{ path: string; matches: LineMatch[] }>;
  /** Head paths the index doesn't hold at that blob (changed since, too large, not text): grep them elsewhere. */
  uncovered: string[];
  limitHit: boolean;
};

export type RepoSearchResult = {
  hits: RepoSearchHit[];
  hasMore: boolean;
  total: number;
  totalCapped: boolean;
};

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: [],
  };
}

/** What the index keeps of a blob: its text when indexable, and its size when it is text at all (languages). */
type ReadBlob = { text: string | null; textBytes: number | null };

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

/**
 * One instance per repository: an FTS5 trigram index of the default branch head in the object's own SQLite,
 * so code search never touches the shared database. Kicks are idempotent; the alarm walks an index job in
 * bounded passes (FILES_PER_PASS blobs each) and abandons a job whose commit is no longer the head.
 */
export class RepoSearch extends DurableObject<Env> {
  private ready = false;

  private get sql(): SqlStorage {
    const sql = this.ctx.storage.sql;
    if (!this.ready) {
      sql.exec(`CREATE TABLE IF NOT EXISTS code_files (
        id INTEGER PRIMARY KEY,
        path TEXT NOT NULL UNIQUE,
        blob_sha TEXT NOT NULL
      )`);
      sql.exec("CREATE VIRTUAL TABLE IF NOT EXISTS code_fts USING fts5(path, content, tokenize = 'trigram')");
      sql.exec("CREATE TABLE IF NOT EXISTS code_sizes (path TEXT PRIMARY KEY, bytes INTEGER NOT NULL)");
      sql.exec(`CREATE TRIGGER IF NOT EXISTS code_files_delete AFTER DELETE ON code_files BEGIN
        DELETE FROM code_fts WHERE rowid = old.id;
      END`);
      this.ready = true;
    }
    return sql;
  }

  /** Records which repository this object indexes and schedules a reconcile with the branch head. */
  async kick(repoId: string): Promise<void> {
    await this.ctx.storage.put("repoId", repoId);
    await this.ctx.storage.setAlarm(Date.now());
  }

  async status(): Promise<{ indexed: RepoSearchStatus; indexing: boolean }> {
    const state = await this.ctx.storage.get<IndexedState>(STATE_KEY);
    const job = await this.ctx.storage.get<Job>(JOB_KEY);
    const indexed = state
      ? { commitSha: state.commitSha, fileCount: state.fileCount, truncated: state.truncated }
      : null;
    return { indexed, indexing: Boolean(job) || (await this.ctx.storage.getAlarm()) !== null };
  }

  /** Kicks a reconcile when `head` isn't the indexed commit; returns whether an index pass is pending. */
  async ensure(repoId: string, head: string): Promise<boolean> {
    const state = await this.ctx.storage.get<IndexedState>(STATE_KEY);
    if (state?.commitSha === head) return (await this.ctx.storage.get<Job>(JOB_KEY)) !== undefined;
    if ((await this.ctx.storage.getAlarm()) === null) await this.kick(repoId);
    return true;
  }

  /** `limit` best files after `offset`, ranked; `terms` are highlighted in the returned lines. */
  async search(
    fts: string,
    terms: string[],
    offset: number,
    limit: number = CODE_PAGE_SIZE,
  ): Promise<RepoSearchResult> {
    const sql = this.sql;
    const rows = sql
      .exec<{ path: string; snippet: string; content: string; rank: number }>(
        `SELECT f.path AS path, ${SNIPPET} AS snippet, code_fts.content AS content, rank
         FROM code_fts JOIN code_files f ON f.id = code_fts.rowid
         WHERE code_fts MATCH ? ORDER BY rank LIMIT ? OFFSET ?`,
        fts,
        limit + 1,
        offset,
      )
      .toArray();
    const counted =
      sql
        .exec<{ total: number }>(
          `SELECT COUNT(*) AS total FROM (SELECT 1 FROM code_fts WHERE code_fts MATCH ? LIMIT ${COUNT_CAP + 1})`,
          fts,
        )
        .one().total ?? 0;
    return {
      hits: rows
        .slice(0, limit)
        .map(({ content, ...row }) => ({ ...row, ...matchingLines(content, terms, LINES_PER_HIT) })),
      hasMore: rows.length > limit,
      total: Math.min(counted, COUNT_CAP),
      totalCapped: counted > COUNT_CAP,
    };
  }

  /**
   * Editor search: line matches for the `head` files ([path, blob sha]) whose indexed copy is that very blob,
   * so any branch close to the default one is mostly searched here, prefiltered by trigram when the query
   * has a literal. Everything else comes back as `uncovered`. Starts indexing a repository never indexed.
   */
  async grep(
    repoId: string,
    query: TextQuery,
    head: Array<[string, string]>,
    limit: number,
    budgetMs: number,
  ): Promise<IndexedGrep> {
    const started = Date.now();
    if (!(await this.ctx.storage.get(STATE_KEY)) && (await this.ctx.storage.getAlarm()) === null)
      await this.kick(repoId);
    const sql = this.sql;
    const indexed = new Map(
      sql
        .exec<{ path: string; blob_sha: string }>("SELECT path, blob_sha FROM code_files")
        .toArray()
        .map((row) => [row.path, row.blob_sha]),
    );
    const covered = new Set<string>();
    const uncovered: string[] = [];
    for (const [path, sha] of head) {
      if (indexed.get(path) === sha) covered.add(path);
      else uncovered.push(path);
    }
    const files: IndexedGrep["files"] = [];
    if (covered.size === 0) return { files, uncovered, limitHit: false };
    const regex = compileQuery(query);
    const literal = requiredLiteral(query);
    const rows = literal
      ? sql.exec<{ path: string; content: string }>(
          `SELECT f.path AS path, code_fts.content AS content FROM code_fts JOIN code_files f ON f.id = code_fts.rowid
           WHERE code_fts MATCH ? ORDER BY f.path`,
          `"${literal.replaceAll('"', '""')}"`,
        )
      : sql.exec<{ path: string; content: string }>(
          `SELECT f.path AS path, code_fts.content AS content FROM code_files f JOIN code_fts ON code_fts.rowid = f.id
           ORDER BY f.path`,
        );
    let found = 0;
    let limitHit = false;
    for (const row of rows) {
      if (!covered.has(row.path)) continue;
      if (found >= limit || Date.now() - started > budgetMs) {
        limitHit = true;
        break;
      }
      const matches = matchLines(row.content, regex, limit - found);
      if (matches.length === 0) continue;
      files.push({ path: row.path, matches });
      found += matches.length;
    }
    return { files, uncovered, limitHit: limitHit || found >= limit };
  }

  /** The repository is gone: drop the whole index. */
  async destroy(): Promise<void> {
    await this.ctx.storage.deleteAlarm();
    await this.ctx.storage.deleteAll();
    this.ready = false;
  }

  override async alarm(): Promise<void> {
    const repoId = await this.ctx.storage.get<string>("repoId");
    if (!repoId) return;
    try {
      if (await this.pass(repoId)) await this.ctx.storage.setAlarm(Date.now());
    } catch (error) {
      console.error("Code index pass failed", repoId, error);
      await this.ctx.storage.setAlarm(Date.now() + RETRY_ON_ERROR_MS);
    }
  }

  /** One bounded step of indexing; returns whether another pass is needed. */
  private async pass(repoId: string): Promise<boolean> {
    const repo = await loadRepoById(repoId);
    if (!repo) {
      await this.destroy();
      return false;
    }
    if (repo.status !== "ready") return false;
    const head = branchSha(await listRefs(repo), repo.default_branch);
    if (!head) return false;
    return withStore(repo.artifacts_name, async (store) => {
      let job = await this.ctx.storage.get<Job>(JOB_KEY);
      if (job && job.commit !== head) job = undefined;
      if (!job) {
        const state = await this.ctx.storage.get<IndexedState>(STATE_KEY);
        const current = state?.schema === INDEX_SCHEMA;
        if (state?.commitSha === head && current) return false;
        const commit = await store.commit(head);
        if (!commit) return false;
        job = {
          commit: head,
          tree: commit.treeSha,
          base: current ? (state?.treeSha ?? null) : null,
          offset: 0,
          schema: INDEX_SCHEMA,
        };
      }
      const plan = await planChanges(store, job.base, job.tree);
      if (job.offset === 0) this.applyDeletes(plan);
      const next = job.offset + FILES_PER_PASS;
      await this.applyBatch(store, plan.upserts.slice(job.offset, next));
      if (next < plan.upserts.length) {
        await this.ctx.storage.put(JOB_KEY, { ...job, offset: next });
        return true;
      }
      await this.finish(job, plan);
      await this.publishLanguages(store, repoId, job.tree);
      // A job planned under an older schema leaves the index stale: rebuild it right away.
      return job.schema !== INDEX_SCHEMA;
    });
  }

  private applyDeletes(plan: Plan): void {
    const sql = this.sql;
    this.ctx.storage.transactionSync(() => {
      if (plan.full) {
        sql.exec("DELETE FROM code_files");
        sql.exec("DELETE FROM code_fts");
        sql.exec("DELETE FROM code_sizes");
      }
      for (const path of new Set(plan.deletes)) {
        sql.exec("DELETE FROM code_files WHERE path = ?", path);
        sql.exec("DELETE FROM code_sizes WHERE path = ?", path);
      }
    });
  }

  private async applyBatch(store: GitStore, items: Plan["upserts"]): Promise<void> {
    const blobs: ReadBlob[] = [];
    for (let index = 0; index < items.length; index += READ_CONCURRENCY) {
      const group = items.slice(index, index + READ_CONCURRENCY);
      blobs.push(...(await Promise.all(group.map((item) => readBlob(store, item.sha)))));
    }
    const sql = this.sql;
    this.ctx.storage.transactionSync(() => {
      for (const [position, item] of items.entries()) {
        const { text, textBytes } = blobs[position] ?? { text: null, textBytes: null };
        if (textBytes === null) sql.exec("DELETE FROM code_sizes WHERE path = ?", item.path);
        else {
          sql.exec(
            `INSERT INTO code_sizes (path, bytes) SELECT ?1, ?2
             WHERE EXISTS (SELECT 1 FROM code_sizes WHERE path = ?1) OR (SELECT COUNT(*) FROM code_sizes) < ?3
             ON CONFLICT (path) DO UPDATE SET bytes = excluded.bytes`,
            item.path,
            textBytes,
            MAX_SIZED_FILES,
          );
        }
        if (text === null) {
          sql.exec("DELETE FROM code_files WHERE path = ?", item.path);
          continue;
        }
        const existing = sql.exec<{ id: number }>("SELECT id FROM code_files WHERE path = ?", item.path).toArray()[0];
        if (!existing) {
          const count = sql.exec<{ total: number }>("SELECT COUNT(*) AS total FROM code_files").one().total;
          if (count >= MAX_INDEXED_FILES) continue;
        }
        sql.exec(
          `INSERT INTO code_files (path, blob_sha) VALUES (?, ?)
           ON CONFLICT (path) DO UPDATE SET blob_sha = excluded.blob_sha`,
          item.path,
          item.sha,
        );
        const id = sql.exec<{ id: number }>("SELECT id FROM code_files WHERE path = ?", item.path).one().id;
        sql.exec("DELETE FROM code_fts WHERE rowid = ?", id);
        sql.exec("INSERT INTO code_fts (rowid, path, content) VALUES (?, ?, ?)", id, item.path, text);
      }
    });
  }

  /** Sums the head's text files per language (root `.gitattributes` applied) into the repository row. */
  private async publishLanguages(store: GitStore, repoId: string, tree: string): Promise<void> {
    const entry = await store.entryAtPath(tree, ".gitattributes");
    const attributes = entry && entry.type !== "tree" && entry.type !== "gitlink" ? await store.blob(entry.sha) : null;
    const rules =
      attributes && attributes.byteLength <= MAX_ATTRIBUTES_BYTES
        ? parseAttributes(new TextDecoder().decode(attributes), MAX_ATTRIBUTE_RULES)
        : [];
    const totals = new Map<string, number>();
    for (const row of this.sql.exec<{ path: string; bytes: number }>("SELECT path, bytes FROM code_sizes")) {
      const language = languageOf(row.path, rules);
      if (language) totals.set(language, (totals.get(language) ?? 0) + row.bytes);
    }
    const languages = summarizeLanguages(totals);
    const encoded = JSON.stringify(languages);
    if ((await this.ctx.storage.get<string>(LANGUAGES_KEY)) === encoded) return;
    const primary = languages.find((language) => language.name !== "Other")?.name ?? null;
    await run("UPDATE repos SET languages = ?, language = ? WHERE id = ?", encoded, primary, repoId);
    await this.ctx.storage.put(LANGUAGES_KEY, encoded);
  }

  private async finish(job: Job, plan: Plan): Promise<void> {
    const total = this.sql.exec<{ total: number }>("SELECT COUNT(*) AS total FROM code_files").one().total;
    const state: IndexedState = {
      commitSha: job.commit,
      treeSha: job.tree,
      fileCount: total,
      truncated: total >= MAX_INDEXED_FILES || (plan.full && plan.upserts.length >= MAX_INDEXED_FILES),
      updatedAt: Date.now(),
      schema: job.schema,
    };
    await this.ctx.storage.put(STATE_KEY, state);
    await this.ctx.storage.delete(JOB_KEY);
  }
}

export function repoSearch(repoId: string) {
  return env.REPO_SEARCH.get(env.REPO_SEARCH.idFromName(repoId));
}