import { pipeline, TextStreamer } from '@huggingface/transformers';

type Device = 'webgpu';
type DType = 'q4f16';
type ChatMessage = { role: 'system' | 'user' | 'assistant'; content: string };

type WorkerRequest =
  | { type: 'load'; modelId: string; revision: string; device: Device; dtype: DType; contextLimit: number }
  | { type: 'generate'; requestId: number; messages: ChatMessage[] };

type ProgressDetails = {
  status: string;
  file?: string;
  progress?: number;
  loaded?: number;
  total?: number;
};

type WorkerResponse =
  | { type: 'progress'; details: ProgressDetails }
  | { type: 'ready'; modelId: string; device: Device; dtype: DType; contextLimit: number }
  | { type: 'prompt-size'; requestId: number; tokens: number; limit: number }
  | { type: 'token'; requestId: number; text: string }
  | { type: 'complete'; requestId: number; text: string; elapsedMs: number }
  | { type: 'too-long'; requestId: number; tokens: number; limit: number }
  | { type: 'error'; stage: 'load' | 'generate'; message: string; requestId?: number };

const MAX_NEW_TOKENS = 192;

const workerScope = self as unknown as {
  addEventListener: (type: 'message', listener: (event: MessageEvent<WorkerRequest>) => void) => void;
  postMessage: (message: WorkerResponse) => void;
};

async function createGenerator(request: Extract<WorkerRequest, { type: 'load' }>) {
  return pipeline('text-generation', request.modelId, {
    revision: request.revision,
    device: request.device,
    dtype: request.dtype,
    progress_callback: (progress) => {
      workerScope.postMessage({
        type: 'progress',
        details: {
          status: progress.status,
          file: 'file' in progress ? progress.file : undefined,
          progress: 'progress' in progress ? progress.progress : undefined,
          loaded: 'loaded' in progress ? progress.loaded : undefined,
          total: 'total' in progress ? progress.total : undefined,
        },
      });
    },
  });
}

type TextGenerator = Awaited<ReturnType<typeof createGenerator>>;

let generator: TextGenerator | null = null;
let activeModelId: string | null = null;
let activeContextLimit = 8192;

async function handleRequest(request: WorkerRequest) {
  if (request.type === 'load') {
    try {
      activeContextLimit = request.contextLimit;
      if (generator && activeModelId !== request.modelId) {
        await generator.dispose();
        generator = null;
      }
      if (!generator) {
        generator = await createGenerator(request);
        activeModelId = request.modelId;
      }
      workerScope.postMessage({
        type: 'ready',
        modelId: request.modelId,
        device: request.device,
        dtype: request.dtype,
        contextLimit: activeContextLimit,
      });
    } catch (error) {
      workerScope.postMessage({
        type: 'error',
        stage: 'load',
        message: error instanceof Error ? error.message : 'The browser could not load the model.',
      });
    }
    return;
  }

  if (!generator) {
    workerScope.postMessage({ type: 'error', stage: 'generate', requestId: request.requestId, message: 'Load the model before sending a question.' });
    return;
  }

  try {
    const maxPromptTokens = activeContextLimit - MAX_NEW_TOKENS;
    const promptTokens = generator.tokenizer.apply_chat_template(request.messages, {
      tokenize: true,
      return_tensor: false,
      return_dict: false,
      add_generation_prompt: true,
    }).length;
    if (promptTokens > maxPromptTokens) {
      workerScope.postMessage({
        type: 'too-long',
        requestId: request.requestId,
        tokens: promptTokens,
        limit: maxPromptTokens,
      });
      return;
    }

    workerScope.postMessage({
      type: 'prompt-size',
      requestId: request.requestId,
      tokens: promptTokens,
      limit: maxPromptTokens,
    });

    const startedAt = performance.now();
    const output = await generator(request.messages, {
      max_new_tokens: MAX_NEW_TOKENS,
      do_sample: false,
      streamer: new TextStreamer(generator.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: true,
        callback_function: (text) => workerScope.postMessage({ type: 'token', requestId: request.requestId, text }),
      }),
    });
    const finalMessage = output[0]?.generated_text.at(-1);
    const text = typeof finalMessage?.content === 'string' ? finalMessage.content : '';
    workerScope.postMessage({
      type: 'complete',
      requestId: request.requestId,
      text,
      elapsedMs: Math.round(performance.now() - startedAt),
    });
  } catch (error) {
    workerScope.postMessage({
      type: 'error',
      stage: 'generate',
      requestId: request.requestId,
      message: error instanceof Error ? error.message : 'The browser could not generate an answer.',
    });
  }
}

workerScope.addEventListener('message', (event) => {
  void handleRequest(event.data);
});