// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license.

/**
 * Thin wrapper over sql.js {@link Database} that mimics the surface of
 * `better-sqlite3`'s `Database` (only the methods the Knex dialect uses).
 *
 * Persistence is handled here, not in the Knex dialect, so it can be
 * orchestrated per-connection without polluting the SQL execution path.
 */

import type { Database, SqlJsStatic } from 'sql.js';

/** When to flush the in-memory WASM database to disk. */
export type PersistMode =
  | 'on-disconnect'
  | 'on-write'
  | `debounced:${number}`;

export interface WasmConnectionOptions {
  /**
   * On-disk file path. `:memory:` (or any value starting with `:`) skips
   * persistence entirely and the database lives only for the process.
   */
  filename: string;
  /** When to persist. Default: `on-disconnect`. */
  persist?: PersistMode;
  /** Pre-loaded sql.js module. If omitted, loaded lazily on first connect. */
  sqlJs?: SqlJsStatic;
  /**
   * Optional override for the `.wasm` locator passed to `initSqlJs()`.
   * Defaults to resolving the file from the `sql.js` package on disk
   * (works in Node and WebContainer).
   */
  locateFile?: (file: string) => string;
  /** Optional logger; defaults to `console`. */
  logger?: { warn: (msg: string, meta?: unknown) => void };
}

/** Mutation method names that should trigger a persistence cycle. */
const WRITE_METHODS = new Set([
  'run',
  'insert',
  'update',
  'del',
  'counter',
]);

/**
 * Detect whether a Node-style `fs` module is available. WebContainer
 * (StackBlitz) provides Node `fs`; pure-browser environments do not.
 */
async function tryLoadFs(): Promise<typeof import('node:fs/promises') | null> {
  try {
    return await import('node:fs/promises');
  } catch {
    return null;
  }
}

/**
 * Resolve a default sql.js WASM locator. We point sql.js at the `.wasm`
 * file shipped inside `sql.js`'s own `dist/` folder. This avoids requiring
 * the caller to host the WASM separately.
 */
async function defaultLocateFile(): Promise<((file: string) => string) | undefined> {
  try {
    const { createRequire } = await import('node:module');
    const require = createRequire(import.meta.url);
    const pkgJsonPath = require.resolve('sql.js/package.json');
    const { dirname, join } = await import('node:path');
    const dir = dirname(pkgJsonPath);
    return (file: string) => join(dir, 'dist', file);
  } catch {
    return undefined;
  }
}

let cachedSqlJs: Promise<SqlJsStatic> | null = null;

async function loadSqlJs(
  locateFile?: (file: string) => string,
): Promise<SqlJsStatic> {
  if (cachedSqlJs) return cachedSqlJs;
  cachedSqlJs = (async () => {
    const mod = await import('sql.js');
    const initSqlJs = (mod as any).default ?? (mod as any);
    const locator = locateFile ?? (await defaultLocateFile());
    const SQL = await initSqlJs(locator ? { locateFile: locator } : undefined);
    return SQL as SqlJsStatic;
  })();
  return cachedSqlJs;
}

/**
 * A sql.js-backed connection that exposes the `prepare`/`exec`/`close`
 * subset used by Knex's SQLite dialect. Mutations are queued through a
 * configurable persistence strategy so the on-disk file stays in sync.
 */
export class WasmSqliteConnection {
  readonly filename: string;
  readonly persist: PersistMode;
  readonly isEphemeral: boolean;

  private db!: Database;
  private fs: typeof import('node:fs/promises') | null = null;
  private dirty = false;
  private debounceMs = 0;
  private debounceTimer: ReturnType<typeof setTimeout> | null = null;
  private flushChain: Promise<void> | null = null;
  private destroyed = false;
  private logger: { warn: (msg: string, meta?: unknown) => void };

  /**
   * Whether a `BEGIN…COMMIT/ROLLBACK` transaction is currently open. Tracked
   * because sql.js's {@link Database.export} closes and reopens the database
   * (it has no in-place serialize), and closing a connection rolls back any
   * open transaction. Flushing mid-transaction would therefore silently
   * abort it, leaving the eventual `COMMIT` to fail with
   * "cannot commit - no transaction is active". We defer the flush until the
   * transaction fully closes. See {@link noteTransactionControl}.
   */
  private rootTxActive = false;
  /** Open `SAVEPOINT` depth (nested transactions emitted by Knex). */
  private savepointDepth = 0;
  /** A flush was requested while a transaction was open; run it on close. */
  private flushDeferred = false;

  /** True while any transaction (root or savepoint) is in flight. */
  private get inTransaction(): boolean {
    return this.rootTxActive || this.savepointDepth > 0;
  }

  constructor(opts: WasmConnectionOptions) {
    this.filename = opts.filename;
    this.persist = opts.persist ?? 'on-disconnect';
    this.isEphemeral =
      this.filename === ':memory:' || this.filename.startsWith(':');
    this.logger = opts.logger ?? console;

    if (typeof this.persist === 'string' && this.persist.startsWith('debounced:')) {
      const ms = Number(this.persist.slice('debounced:'.length));
      this.debounceMs = Number.isFinite(ms) && ms > 0 ? ms : 250;
    }
  }

  /** Open the underlying sql.js database, loading bytes from disk if any. */
  async open(sqlJs?: SqlJsStatic, locateFile?: (file: string) => string): Promise<void> {
    const SQL = sqlJs ?? (await loadSqlJs(locateFile));

    if (this.isEphemeral) {
      this.db = new SQL.Database();
      return;
    }

    this.fs = await tryLoadFs();
    if (!this.fs) {
      this.logger.warn(
        '[driver-sqlite-wasm] No node:fs available — falling back to in-memory database. ' +
          'Data will not be persisted across reloads.',
      );
      this.db = new SQL.Database();
      return;
    }

    // Ensure parent directory exists, then load bytes if the file exists.
    const { dirname } = await import('node:path');
    const dir = dirname(this.filename);
    if (dir && dir !== '.') {
      await this.fs.mkdir(dir, { recursive: true });
    }

    let bytes: Uint8Array | undefined;
    try {
      const buf = await this.fs.readFile(this.filename);
      bytes = new Uint8Array(buf.buffer, buf.byteOffset, buf.byteLength);
    } catch (e: any) {
      if (e?.code !== 'ENOENT') throw e;
    }

    this.db = bytes ? new SQL.Database(bytes) : new SQL.Database();
  }

  /**
   * Update transaction state from a transaction-control statement and, when a
   * transaction has just fully closed, run any flush that was deferred while
   * it was open. Called by the Knex dialect for every `BEGIN` / `COMMIT` /
   * `ROLLBACK` / `SAVEPOINT` / `RELEASE` statement.
   *
   * We bias toward "in transaction": an unrecognised form leaves the flag set,
   * which at worst delays a flush (safe) rather than exporting mid-transaction
   * (which would abort it).
   */
  noteTransactionControl(sql: string): void {
    const s = sql.trim().toUpperCase();
    if (/^BEGIN\b/.test(s)) {
      this.rootTxActive = true;
    } else if (/^(COMMIT|END)\b/.test(s)) {
      // A COMMIT/END ends the whole transaction regardless of savepoint nesting.
      this.rootTxActive = false;
      this.savepointDepth = 0;
    } else if (/^ROLLBACK\s+TO\b/.test(s)) {
      // Rolls back to a savepoint but keeps the (outer) transaction open.
    } else if (/^ROLLBACK\b/.test(s)) {
      this.rootTxActive = false;
      this.savepointDepth = 0;
    } else if (/^SAVEPOINT\b/.test(s)) {
      this.savepointDepth += 1;
    } else if (/^RELEASE\b/.test(s)) {
      this.savepointDepth = Math.max(0, this.savepointDepth - 1);
    }
    // If the transaction just fully closed and a flush was deferred while it
    // was open, run it now. We key off `flushDeferred` (set only when
    // `markDirty` actually wanted to flush) rather than `dirty`, so persist
    // modes that don't flush per-write — e.g. `on-disconnect` — still defer to
    // close() instead of flushing on every COMMIT.
    if (!this.inTransaction && this.flushDeferred) {
      this.flushDeferred = false;
      void this.flush();
    }
  }

  /** Hint that a mutation just executed; schedule a flush if needed. */
  markDirty(method?: string): void {
    if (this.isEphemeral || !this.fs) return;
    if (method && !WRITE_METHODS.has(method)) return;
    this.dirty = true;

    if (this.persist === 'on-write') {
      void this.flush();
      return;
    }
    if (this.debounceMs > 0) {
      if (this.debounceTimer) clearTimeout(this.debounceTimer);
      this.debounceTimer = setTimeout(() => {
        this.debounceTimer = null;
        void this.flush();
      }, this.debounceMs);
    }
    // 'on-disconnect' → flush only at close()
  }

  /**
   * Force a write of the current database state to disk.
   *
   * Flushes are strictly serialized through a single promise chain: every call
   * appends an export+write step that runs after all previously-queued steps.
   * This matters because sql.js `export()` mutates the live connection (it
   * closes and reopens the database), so two exports must never overlap — and
   * because the returned promise must not resolve until the caller's own write
   * has hit disk (deterministic for tests and for `close()`). Each step
   * re-checks `dirty` at run time, so a no-op write collapses cheaply and a
   * write that arrived mid-flush is captured by the next queued step.
   */
  async flush(): Promise<void> {
    if (this.isEphemeral || !this.fs || this.destroyed) return;
    // Never export while a transaction is open: sql.js's `export()` closes and
    // reopens the database, which rolls back the in-flight transaction and
    // makes the subsequent COMMIT fail. Defer until the transaction closes
    // (handled in `noteTransactionControl`).
    if (this.inTransaction) {
      this.flushDeferred = true;
      return;
    }

    const prev = this.flushChain;
    const step = (prev ?? Promise.resolve()).then(async () => {
      if (!this.dirty || this.destroyed || this.inTransaction) return;
      // Snapshot dirty=false before export so a concurrent write re-marks us
      // and is picked up by the next queued step.
      this.dirty = false;
      try {
        const exported = this.db.export();
        // sql.js returns a Uint8Array; Buffer.from on it shares memory but
        // works fine for the synchronous writeFile enqueue below.
        await this.fs!.writeFile(this.filename, Buffer.from(exported));
      } catch (err) {
        this.dirty = true; // let a later flush retry
        throw err;
      }
    });
    // Keep the chain tail alive but swallow its rejection there so one failed
    // flush doesn't poison every future flush; the awaited `step` still throws.
    this.flushChain = step.catch(() => {});
    await step;
  }

  /** Close the database, flushing any pending writes first. */
  async close(): Promise<void> {
    if (this.destroyed) return;
    if (this.debounceTimer) {
      clearTimeout(this.debounceTimer);
      this.debounceTimer = null;
    }
    // Any transaction still open at close is abandoned and will be rolled back
    // by `db.close()`; clear the flag so the final flush is not deferred and
    // already-committed data is persisted.
    this.rootTxActive = false;
    this.savepointDepth = 0;
    try {
      await this.flush();
    } finally {
      this.destroyed = true;
      try {
        this.db.close();
      } catch {
        /* ignore */
      }
    }
  }

  /** Access the raw sql.js database (for the Knex dialect). */
  get raw(): Database {
    return this.db;
  }
}
