/**
 * postMessage protocol between the UI thread and the inference worker.
 */

import type { AlgorithmId, Detection, GenerationResult, SamplingConfig, SentenceTrace } from '../watermark/types';
import type { KirchenbauerParams } from '../watermark/kirchenbauer';
import type { KsemstampParams } from '../watermark/ksemstamp';
import type { TextsealParams } from '../watermark/textseal';

export interface AlgorithmParams {
  kirchenbauer: KirchenbauerParams;
  ksemstamp: KsemstampParams;
  textseal: TextsealParams;
}

/** Top-N next-token candidates for the wizard's Stage 1 / Stage 2 views. */
export interface NextTokenCandidate {
  tokenId: number;
  text: string;
  logit: number;
  prob: number;
}

export interface DistributionResult {
  /** Last token of the templated prompt - the Kirchenbauer PRF context. */
  prevTokenId: number;
  prevTokenText: string;
  /** Tail of the prompt (oldest first) so wider PRF contexts can be shown. */
  prevTokenIds: number[];
  vocabSize: number;
  candidates: NextTokenCandidate[];
}

/** UI -> worker */
export type WorkerRequest =
  | { type: 'check' }
  | { type: 'load'; modelId: string }
  | { type: 'distribution'; requestId: number; prompt: string; topN: number; temperature: number }
  | {
      type: 'generate';
      requestId: number;
      algorithm: AlgorithmId;
      prompt: string;
      sampling: SamplingConfig;
      params: AlgorithmParams;
      /** Per-algorithm secret entered by the user in Stage 1 (decimal string). */
      secretKey?: string;
    }
  | {
      type: 'detect';
      requestId: number;
      text: string;
      sampling: SamplingConfig;
      params: AlgorithmParams;
      /** Per-algorithm secrets so all three detectors can score the same text. */
      keys?: { kirchenbauer?: string; ksemstamp?: string; textseal?: string };
      /** Run the TextSeal entropy-weighted variant (needs a model forward). */
      withEntropy?: boolean;
    };

export interface DeviceInfo {
  webgpuSupported: boolean;
  shaderF16: boolean;
  device: 'webgpu' | 'wasm';
  dtype: string;
  adapterInfo?: { vendor?: string; architecture?: string };
}

export interface DetectionBundle {
  kirchenbauer: Detection;
  textseal: Detection;
  ksemstamp: Detection | null;
}

/** worker -> UI */
export type WorkerResponse =
  | { type: 'check-result'; info: DeviceInfo }
  | { type: 'load-progress'; file: string; progress: number; loadedMB: number; totalMB: number }
  | { type: 'load-done'; modelId: string; initMs: number; info: DeviceInfo }
  | { type: 'load-error'; error: string }
  | { type: 'distribution-done'; requestId: number; result: DistributionResult }
  | { type: 'distribution-error'; requestId: number; error: string }
  | {
      type: 'generate-progress';
      requestId: number;
      algorithm: AlgorithmId;
      text: string;
      tokensDone: number;
      retries: number;
    }
  /**
   * One sampled token, streamed as it is chosen. `index` is authoritative:
   * after a k-SemStamp rollback it goes backwards, and the receiver should
   * drop everything from that index on before appending.
   */
  | {
      type: 'token';
      requestId: number;
      algorithm: AlgorithmId;
      index: number;
      text: string;
      /** Kirchenbauer: which list the sampled token belonged to. */
      green?: boolean;
      /** TextSeal: which of the two keys routed this step. */
      keyId?: 1 | 2;
      entropy?: number;
    }
  /** k-SemStamp candidate sentence judged (streamed live during rejection sampling). */
  | { type: 'candidate'; requestId: number; algorithm: AlgorithmId; sentence: SentenceTrace }
  | { type: 'generate-done'; requestId: number; result: GenerationResult }
  | { type: 'generate-error'; requestId: number; algorithm: AlgorithmId; error: string }
  | { type: 'detect-done'; requestId: number; detections: DetectionBundle }
  | { type: 'detect-error'; requestId: number; error: string };