import {
  EMBEDDING_MODEL_ID,
  type EmbeddingCapabilities,
  type EmbeddingEngine,
  type EmbeddingModelState,
  type EmbeddingPlatform,
  type ProjectMeaningStore,
  type StoredProjectMeaningIndex,
} from '@learnitive/project-meaning';

const DB_NAME = 'interactive-project-meaning-v1';
const DB_VERSION = 1;
const INDEX_STORE = 'indices';

type WorkerRequest = {
  id: string;
  type: 'capabilities' | 'model-state' | 'prepare-model' | 'remove-model' | 'load' | 'unload' | 'embed-documents' | 'embed-query' | 'cancel';
  texts?: readonly string[];
  query?: string;
  allowDownload?: boolean;
};

type WorkerResponse =
  | { id: string; type: 'response'; ok: true; value?: unknown }
  | { id: string; type: 'response'; ok: false; error: string }
  | { id: string; type: 'progress'; state: EmbeddingModelState };

type Pending = {
  resolve(value: unknown): void;
  reject(error: Error): void;
  onProgress?: (state: EmbeddingModelState) => void;
};

function openDatabase(): Promise<IDBDatabase> {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(DB_NAME, DB_VERSION);
    request.onupgradeneeded = () => {
      const db = request.result;
      if (!db.objectStoreNames.contains(INDEX_STORE)) db.createObjectStore(INDEX_STORE);
    };
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error ?? new Error('Unable to open project meaning database'));
  });
}

async function withStore<T>(mode: IDBTransactionMode, fn: (store: IDBObjectStore) => IDBRequest<T>): Promise<T> {
  const db = await openDatabase();
  try {
    return await new Promise<T>((resolve, reject) => {
      const transaction = db.transaction(INDEX_STORE, mode);
      const request = fn(transaction.objectStore(INDEX_STORE));
      request.onsuccess = () => resolve(request.result);
      request.onerror = () => reject(request.error ?? new Error('Project meaning storage operation failed'));
      transaction.onabort = () => reject(transaction.error ?? new Error('Project meaning storage transaction aborted'));
    });
  } finally {
    db.close();
  }
}

class IndexedDbProjectMeaningStore implements ProjectMeaningStore {
  async load(memoryId: string): Promise<StoredProjectMeaningIndex | undefined> {
    return withStore('readonly', (store) => store.get(memoryId));
  }

  async save(memoryId: string, index: StoredProjectMeaningIndex): Promise<void> {
    await withStore('readwrite', (store) => store.put(index, memoryId));
  }

  async clear(memoryId: string): Promise<void> {
    await withStore('readwrite', (store) => store.delete(memoryId));
  }
}

class EmbeddingWorkerClient {
  private worker?: Worker;
  private readonly pending = new Map<string, Pending>();

  private ensureWorker(): Worker {
    if (this.worker) return this.worker;
    const worker = new Worker(new URL('./embedding.worker.ts', import.meta.url), {
      type: 'module',
      name: 'interactive-embedding-gemma',
    });
    worker.addEventListener('message', (event: MessageEvent<WorkerResponse>) => {
      const message = event.data;
      if (message.type === 'progress') {
        this.pending.get(message.id)?.onProgress?.(message.state);
        return;
      }
      const pending = this.pending.get(message.id);
      if (!pending) return;
      this.pending.delete(message.id);
      if (message.ok) pending.resolve(message.value);
      else pending.reject(new Error(message.error));
    });
    worker.addEventListener('error', (event) => {
      const error = new Error(event.message || 'Embedding worker failed');
      for (const pending of this.pending.values()) pending.reject(error);
      this.pending.clear();
      worker.terminate();
      if (this.worker === worker) this.worker = undefined;
    });
    this.worker = worker;
    return worker;
  }

  call<T>(
    request: Omit<WorkerRequest, 'id'>,
    options: { signal?: AbortSignal; onProgress?: (state: EmbeddingModelState) => void } = {},
  ): Promise<T> {
    const id = crypto.randomUUID();
    const worker = this.ensureWorker();
    return new Promise<T>((resolve, reject) => {
      const onAbort = () => {
        worker.postMessage({ id, type: 'cancel' } satisfies WorkerRequest);
        this.pending.delete(id);
        reject(new DOMException('The operation was aborted', 'AbortError'));
      };
      if (options.signal?.aborted) {
        onAbort();
        return;
      }
      options.signal?.addEventListener('abort', onAbort, { once: true });
      this.pending.set(id, {
        resolve: (value) => {
          options.signal?.removeEventListener('abort', onAbort);
          resolve(value as T);
        },
        reject: (error) => {
          options.signal?.removeEventListener('abort', onAbort);
          reject(error);
        },
        onProgress: options.onProgress,
      });
      worker.postMessage({ ...request, id } satisfies WorkerRequest);
    });
  }

  dispose(): void {
    const error = new Error('Embedding worker disposed');
    for (const pending of this.pending.values()) pending.reject(error);
    this.pending.clear();
    this.worker?.terminate();
    this.worker = undefined;
  }
}

class BrowserEmbeddingEngine implements EmbeddingEngine {
  constructor(private readonly client: EmbeddingWorkerClient) {}

  async load(signal?: AbortSignal): Promise<void> {
    await this.client.call<void>({ type: 'load', allowDownload: false }, { signal });
  }

  async unload(): Promise<void> {
    await this.client.call<void>({ type: 'unload' }).catch(() => undefined);
  }

  async embedDocuments(texts: readonly string[], signal?: AbortSignal): Promise<Float32Array[]> {
    if (texts.length === 0) return [];
    return this.client.call<Float32Array[]>({ type: 'embed-documents', texts }, { signal });
  }

  async embedQuery(query: string, signal?: AbortSignal): Promise<Float32Array> {
    return this.client.call<Float32Array>({ type: 'embed-query', query }, { signal });
  }
}

export function createBrowserEmbeddingPlatform(): EmbeddingPlatform {
  const client = new EmbeddingWorkerClient();
  return {
    modelId: EMBEDDING_MODEL_ID,
    capabilities: () => client.call<EmbeddingCapabilities>({ type: 'capabilities' }),
    modelState: () => client.call<EmbeddingModelState>({ type: 'model-state' }),
    prepareModel: (onProgress, signal) => client.call<EmbeddingModelState>({ type: 'prepare-model', allowDownload: true }, { signal, onProgress }),
    removeModel: () => client.call<void>({ type: 'remove-model' }),
    createEngine: () => new BrowserEmbeddingEngine(client),
    store: new IndexedDbProjectMeaningStore(),
    dispose: () => client.dispose(),
  };
}