import { CONTEXT_ID, PROTOCOL_SCHEMA, isEnvelope, type MessageEnvelope } from "./protocol";

type Listener = (message: MessageEnvelope) => void;

function commandError(message: MessageEnvelope): Error {
  const payload = message.payload && typeof message.payload === "object" ? message.payload as Record<string, unknown> : {};
  const error = new Error(String(payload.reason ?? payload.gate_reason ?? payload.message_key ?? payload.code ?? "worker_error"));
  Object.assign(error, {
    code: String(payload.code ?? "worker_error"),
    reason: payload.reason ?? payload.gate_reason ?? null,
    message_key: payload.message_key ?? null,
    response: message,
  });
  return error;
}

export class ReactWorkerClient {
  private readonly worker: Worker;
  private readonly contextId: string;
  private readonly sessionId: string;
  private sequence = 0;
  private uiSequence = 2000000000;
  private listeners = new Set<Listener>();
  private pending = new Map<number, { resolve: (message: MessageEnvelope) => void; reject: (error: Error) => void }>();
  private readyPromise: Promise<void> | null = null;

  constructor({ worker, contextId = CONTEXT_ID, sessionId = `react-${Date.now()}` }: { worker?: Worker; contextId?: string; sessionId?: string } = {}) {
    // The entry imports the source-compatible SimSession; keep it an ES module
    // so Vite does not wrap it as a classic worker and drop its imports.
    this.worker = worker ?? new Worker(new URL("./sim-worker-entry.ts", import.meta.url), { type: "module" });
    this.contextId = contextId;
    this.sessionId = sessionId;
    this.worker.onmessage = (event: MessageEvent<unknown>) => {
      if (!isEnvelope(event.data)) return;
      const uiStateResponse = event.data.kind === "UI_STATE_RESULT";
      if (!uiStateResponse) for (const listener of this.listeners) listener(event.data);
      const pending = this.pending.get(event.data.seq);
      if (pending) {
        this.pending.delete(event.data.seq);
        if (event.data.kind === "ERROR" || event.data.payload && typeof event.data.payload === "object" && (event.data.payload as { status?: string }).status === "blocked") pending.reject(commandError(event.data));
        else pending.resolve(event.data);
      }
    };
    this.worker.onerror = (event) => {
      const message: MessageEnvelope = {
        schema_version: PROTOCOL_SCHEMA,
        kind: "ERROR",
        context_id: this.contextId,
        session_id: this.sessionId,
        seq: ++this.sequence,
        payload: { code: "worker_transport_error", message_key: event.message || "worker_failed_to_load" },
      };
      for (const listener of this.listeners) listener(message);
      for (const pending of this.pending.values()) pending.reject(new Error("worker_transport_error"));
      this.pending.clear();
    };
  }

  subscribe(listener: Listener): () => void { this.listeners.add(listener); return () => this.listeners.delete(listener); }

  send(kind: string, payload: unknown = {}, operationId?: string): Promise<MessageEnvelope> {
    const message: MessageEnvelope = { schema_version: PROTOCOL_SCHEMA, kind, context_id: this.contextId, session_id: this.sessionId, seq: ++this.sequence, ...(operationId ? { operation_id: operationId } : {}), payload };
    return new Promise((resolve, reject) => { this.pending.set(message.seq, { resolve, reject }); this.worker.postMessage(message); });
  }

  private sendUiState(kind: "UI_STATE_SAVE" | "UI_STATE_LOAD" | "UI_STATE_CLEAR", payload: Record<string, unknown>): Promise<MessageEnvelope> {
    const seq = ++this.uiSequence;
    const message: MessageEnvelope = { schema_version: PROTOCOL_SCHEMA, kind, context_id: this.contextId, session_id: this.sessionId, seq, operation_id: `ui-cache-${seq}`, payload };
    return new Promise((resolve, reject) => { this.pending.set(seq, { resolve, reject }); this.worker.postMessage(message); });
  }

  saveUiState(state: Record<string, unknown>, key = "xyzac-trt"): Promise<MessageEnvelope> { return this.sendUiState("UI_STATE_SAVE", { key, state }); }
  loadUiState(key = "xyzac-trt"): Promise<MessageEnvelope> { return this.sendUiState("UI_STATE_LOAD", { key }); }
  clearUiState(key = "xyzac-trt"): Promise<MessageEnvelope> { return this.sendUiState("UI_STATE_CLEAR", { key }); }

  hello(): Promise<MessageEnvelope> { return this.send("HELLO", { supported_schema_versions: [PROTOCOL_SCHEMA], requested_features: ["snapshot", "trajectory"] }); }
  attach(): Promise<MessageEnvelope> { return this.send("SESSION_ATTACH", {}, "react-session-attach"); }
  ensureSession(): Promise<void> {
    if (!this.readyPromise) this.readyPromise = this.hello().then(() => this.attach()).then(() => undefined);
    return this.readyPromise;
  }
  close(): void { this.worker.terminate(); }
}