import {
  LOCAL_AI_MODELS,
  type LocalAiCapabilities,
  type LocalAiModelId,
  type LocalAiModelState,
  type LocalAiPlatform,
} from '@learnitive/project-ai';
import { createBrowserLiteRtGenerationEngine } from './litert-main-runtime';

type WorkerRequest = {
  id: string;
  type: 'capabilities' | 'model-state' | 'download' | 'remove' | 'cancel';
  modelId?: LocalAiModelId;
};

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

type PendingRequest = {
  resolve: (value: unknown) => void;
  reject: (reason: Error) => void;
  onProgress?: (state: LocalAiModelState) => void;
};

class LocalAiWorkerClient {
  private worker: Worker | undefined;
  private readonly pending = new Map<string, PendingRequest>();

  private ensureWorker(): Worker {
    if (this.worker) return this.worker;
    const worker = new Worker(new URL('./local-ai.worker.ts', import.meta.url), {
      type: 'module',
      name: 'interactive-local-ai',
    });
    worker.addEventListener('message', (event: MessageEvent<WorkerResponse>) => this.onMessage(event.data));
    worker.addEventListener('error', (event) => {
      const error = new Error(event.message || 'Local AI 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;
  }

  private onMessage(message: WorkerResponse): void {
    if (message.type === 'progress') {
      this.pending.get(message.id)?.onProgress?.(message.state);
      return;
    }
    const pending = this.pending.get(message.id);
    if (pending) {
      this.pending.delete(message.id);
      if (message.ok) pending.resolve(message.value);
      else pending.reject(new Error(message.error));
    }
  }

  call<T>(
    request: Omit<WorkerRequest, 'id'>,
    options: { signal?: AbortSignal; onProgress?: (state: LocalAiModelState) => 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: (reason) => {
          options.signal?.removeEventListener('abort', onAbort);
          reject(reason);
        },
        onProgress: options.onProgress,
      });
      worker.postMessage({ ...request, id } satisfies WorkerRequest);
    });
  }

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


export function createBrowserLocalAiPlatform(): LocalAiPlatform & { dispose(): void } {
  const client = new LocalAiWorkerClient();
  return {
    models: LOCAL_AI_MODELS,
    capabilities: () => client.call<LocalAiCapabilities>({ type: 'capabilities' }),
    modelState: (modelId) => client.call<LocalAiModelState>({ type: 'model-state', modelId }),
    downloadModel: (modelId, onProgress, signal) =>
      client.call<LocalAiModelState>({ type: 'download', modelId }, { signal, onProgress }),
    removeModel: (modelId) => client.call<void>({ type: 'remove', modelId }),
    createGenerationEngine: (modelId) => createBrowserLiteRtGenerationEngine(modelId),
    dispose: () => client.dispose(),
  };
}