import type {
  WorkerRequest,
  WorkerRequestPayload,
  WorkerResponse,
  WorkerSuccessPayload,
} from "./protocol";
import type {
  LoadProgress,
  TokenizationResult,
  TokenizerMetadata,
} from "./types";

interface PendingRequest {
  resolve: (payload: WorkerSuccessPayload) => void;
  reject: (error: Error) => void;
  onProgress?: (progress: LoadProgress) => void;
}

export class TokenizerClient {
  private readonly worker = new Worker(
    new URL("./tokenizer.worker.ts", import.meta.url),
    { type: "module" },
  );
  private readonly pending = new Map<number, PendingRequest>();
  private nextId = 1;

  constructor() {
    this.worker.onmessage = (event: MessageEvent<WorkerResponse>) => {
      const message = event.data;
      const request = this.pending.get(message.id);
      if (!request) return;

      if (message.type === "progress") {
        request.onProgress?.(message.progress);
        return;
      }

      this.pending.delete(message.id);
      if (message.type === "error") {
        request.reject(new Error(message.message));
      } else {
        request.resolve(message.payload);
      }
    };

    this.worker.onerror = (event) => {
      const error = new Error(event.message || "Tokenizer Worker 启动失败。");
      this.pending.forEach((request) => request.reject(error));
      this.pending.clear();
    };
  }

  private request(
    request: WorkerRequestPayload,
    onProgress?: (progress: LoadProgress) => void,
  ): Promise<WorkerSuccessPayload> {
    const id = this.nextId++;
    return new Promise((resolve, reject) => {
      this.pending.set(id, { resolve, reject, onProgress });
      this.worker.postMessage({ ...request, id } as WorkerRequest);
    });
  }

  async load(
    reference: string,
    onProgress: (progress: LoadProgress) => void,
  ): Promise<TokenizerMetadata> {
    return (await this.request({ type: "load", reference }, onProgress)) as TokenizerMetadata;
  }

  async tokenize(
    text: string,
    addSpecialTokens: boolean,
  ): Promise<TokenizationResult> {
    return (await this.request({ type: "tokenize", text, addSpecialTokens })) as TokenizationResult;
  }

  async cacheSize(): Promise<number> {
    return (await this.request({ type: "cache-size" })) as number;
  }

  async clearCache(): Promise<void> {
    await this.request({ type: "clear-cache" });
  }

  dispose(): void {
    this.worker.terminate();
    const error = new Error("Tokenizer Worker 已关闭。");
    this.pending.forEach((request) => request.reject(error));
    this.pending.clear();
  }
}