import {
  AsyncDuckDB,
  ConsoleLogger,
  DuckDBDataProtocol,
  getJsDelivrBundles,
  selectBundle,
} from '@duckdb/duckdb-wasm';

export type QueryMonthFile = {
  month: string;
  path: string;
};

type QueryEngine = {
  database: AsyncDuckDB;
  connection: Awaited<ReturnType<AsyncDuckDB['connect']>>;
  worker: Worker;
  registeredFiles: Map<string, string>;
};

let queryEnginePromise: Promise<QueryEngine> | null = null;

async function createQueryEngine(): Promise<QueryEngine> {
  const bundle = await selectBundle(getJsDelivrBundles());
  const workerUrl = URL.createObjectURL(
    new Blob([`importScripts("${bundle.mainWorker}");`], {
      type: 'text/javascript',
    }),
  );
  const worker = new Worker(workerUrl);
  const database = new AsyncDuckDB(new ConsoleLogger(), worker);

  try {
    await database.instantiate(bundle.mainModule, bundle.pthreadWorker);
    URL.revokeObjectURL(workerUrl);
    const connection = await database.connect();
    const engine = {
      database,
      connection,
      worker,
      registeredFiles: new Map<string, string>(),
    };

    window.addEventListener('pagehide', () => {
      void connection.close().catch(() => undefined);
      void database.terminate().catch(() => undefined);
      worker.terminate();
    }, { once: true });

    return engine;
  } catch (error) {
    URL.revokeObjectURL(workerUrl);
    worker.terminate();
    await database.terminate().catch(() => undefined);
    throw error;
  }
}

export function getQueryEngine(): Promise<QueryEngine> {
  if (!queryEnginePromise) {
    queryEnginePromise = createQueryEngine().catch((error: unknown) => {
      queryEnginePromise = null;
      throw error;
    });
  }
  return queryEnginePromise;
}

export async function registerMonthFiles(
  engine: QueryEngine,
  months: QueryMonthFile[],
): Promise<string[]> {
  return Promise.all(months.map(async (month) => {
    const filename = `${month.month}.slim.parquet`;
    const url = new URL(month.path, window.location.href).toString();
    const registeredUrl = engine.registeredFiles.get(filename);

    if (registeredUrl === url) return filename;
    if (registeredUrl) {
      throw new Error(`A different snapshot file is already registered as ${filename}.`);
    }

    await engine.database.registerFileURL(
      filename,
      url,
      DuckDBDataProtocol.HTTP,
      false,
    );
    engine.registeredFiles.set(filename, url);
    return filename;
  }));
}