import { TransferAbortError, TransferTransportError } from '../errors.js';
import type {
  ChunkAck,
  ChunkSendOptions,
  ITransferTransport,
  TransferResumeState,
} from './transport.js';
import type { TransferAuthenticator } from './http-transport.js';
import { NoopAuthenticator } from './http-transport.js';

/**
 * WebSocket-based chunk transport (opt-in).
 *
 * One connection per peer, multiplexed across all in-flight streams. Each
 * outgoing chunk is sent as a binary frame; the server replies with a JSON
 * ACK keyed by an internal request id. Failed connections trigger
 * `TransferTransportError`; the wrapping `FallbackTransferTransport` swaps
 * to HTTP.
 *
 * Wire format (client→server):
 *   1 byte type (0x01 = chunk, 0x02 = resume-state-query)
 *   16 bytes requestId (opaque)
 *   payload (chunk: streamId(16) + laneId(4 BE) + seq(8 BE) + envelopeBytes;
 *            resume-query: streamId(16))
 *
 * Wire format (server→client):
 *   1 byte type (0x81 = chunk-ack, 0x82 = resume-state, 0xFE = error)
 *   16 bytes requestId (echoes the request)
 *   payload (chunk-ack: u32 lastSeq + u32 bytesReceived;
 *            resume-state: JSON bytes;
 *            error: JSON bytes)
 */
export interface ShadeTransferWsTransportOptions {
  resolveWsUrl: (peerAddress: string) => Promise<string>;
  authenticator?: TransferAuthenticator;
  /** WebSocket constructor override (browsers vs node). */
  WebSocketCtor?: typeof WebSocket;
  /** Connect timeout in ms. */
  connectTimeoutMs?: number;
}

const TYPE_CHUNK = 0x01;
const TYPE_RESUME_QUERY = 0x02;
const TYPE_CHUNK_ACK = 0x81;
const TYPE_RESUME_STATE = 0x82;
const TYPE_ERROR = 0xfe;

export class ShadeTransferWsTransport implements ITransferTransport {
  private readonly auth: TransferAuthenticator;
  private readonly WebSocketCtor: typeof WebSocket;
  private readonly connectTimeoutMs: number;
  private readonly connections = new Map<string, Promise<WebSocket>>();

  constructor(private readonly options: ShadeTransferWsTransportOptions) {
    this.auth = options.authenticator ?? NoopAuthenticator;
    const ctor =
      options.WebSocketCtor ??
      ((globalThis as unknown as { WebSocket?: typeof WebSocket }).WebSocket);
    if (ctor === undefined) {
      throw new TransferTransportError('WebSocket constructor not available');
    }
    this.WebSocketCtor = ctor;
    this.connectTimeoutMs = options.connectTimeoutMs ?? 5000;
  }

  async probe(peerAddress: string): Promise<void> {
    // Probe by opening a connection; close immediately on success.
    const ws = await this.connect(peerAddress);
    void ws;
  }

  async sendChunk(
    peerAddress: string,
    streamId: string,
    laneId: number,
    seq: number | bigint,
    bytes: Uint8Array,
    options?: ChunkSendOptions,
  ): Promise<ChunkAck> {
    const ws = await this.connect(peerAddress);
    const requestId = randomBytes(16);
    const seqBig = typeof seq === 'bigint' ? seq : BigInt(seq);

    const sidBytes = base64UrlToBytes(streamId);
    if (sidBytes.length !== 16) {
      throw new TransferTransportError(`streamId must decode to 16 bytes`);
    }
    const header = new Uint8Array(1 + 16 + 16 + 4 + 8);
    header[0] = TYPE_CHUNK;
    header.set(requestId, 1);
    header.set(sidBytes, 17);
    new DataView(header.buffer).setUint32(33, laneId, false);
    new DataView(header.buffer).setBigUint64(37, seqBig, false);
    const frame = new Uint8Array(header.length + bytes.length);
    frame.set(header, 0);
    frame.set(bytes, header.length);

    // Sign chunk for auth — included in WS-level metadata frame would be
    // ideal, but for v0.2.0 we just call the authenticator and discard,
    // matching the HTTP path's contract. Real auth lands in 0.3.
    void (await this.auth.signChunk({
      streamId,
      laneId,
      seq: seqBig,
      bodyHash: await sha256(bytes),
    }));

    const response = await this.request(ws, requestId, frame, options?.signal);
    const view = new DataView(response);
    const type = view.getUint8(0);
    if (type === TYPE_ERROR) {
      const text = new TextDecoder().decode(new Uint8Array(response, 17));
      throw new TransferTransportError(`WS sendChunk error: ${text}`);
    }
    if (type !== TYPE_CHUNK_ACK) {
      throw new TransferTransportError(`unexpected WS response type 0x${type.toString(16)}`);
    }
    const lastSeq = view.getUint32(17, false);
    const bytesReceived = view.getUint32(21, false);
    return { lastSeq, bytesReceived };
  }

  async fetchResumeState(
    peerAddress: string,
    streamId: string,
  ): Promise<TransferResumeState | null> {
    const ws = await this.connect(peerAddress);
    const requestId = randomBytes(16);
    const sidBytes = base64UrlToBytes(streamId);
    const frame = new Uint8Array(1 + 16 + 16);
    frame[0] = TYPE_RESUME_QUERY;
    frame.set(requestId, 1);
    frame.set(sidBytes, 17);

    const response = await this.request(ws, requestId, frame);
    const view = new DataView(response);
    const type = view.getUint8(0);
    if (type === TYPE_ERROR) {
      const text = new TextDecoder().decode(new Uint8Array(response, 17));
      if (text.includes('not found')) return null;
      throw new TransferTransportError(`WS fetchResumeState error: ${text}`);
    }
    if (type !== TYPE_RESUME_STATE) {
      throw new TransferTransportError(`unexpected WS response type 0x${type.toString(16)}`);
    }
    const json = new TextDecoder().decode(new Uint8Array(response, 17));
    return JSON.parse(json) as TransferResumeState;
  }

  /** Force-close all connections (call from `engine.destroy()`). */
  closeAll(): void {
    for (const conn of this.connections.values()) {
      void conn.then((ws) => ws.close()).catch(() => {});
    }
    this.connections.clear();
  }

  private connect(peerAddress: string): Promise<WebSocket> {
    let conn = this.connections.get(peerAddress);
    if (conn === undefined) {
      conn = this.openConnection(peerAddress);
      this.connections.set(peerAddress, conn);
      conn.catch(() => {
        this.connections.delete(peerAddress);
      });
    }
    return conn;
  }

  private async openConnection(peerAddress: string): Promise<WebSocket> {
    const url = await this.options.resolveWsUrl(peerAddress);
    const ws = new this.WebSocketCtor(url);
    ws.binaryType = 'arraybuffer';

    return new Promise<WebSocket>((resolve, reject) => {
      const timer = setTimeout(() => {
        ws.close();
        reject(new TransferTransportError(`WS connect timeout to ${url}`));
      }, this.connectTimeoutMs);

      ws.addEventListener('open', () => {
        clearTimeout(timer);
        resolve(ws);
      });
      ws.addEventListener('error', () => {
        clearTimeout(timer);
        reject(new TransferTransportError(`WS connect failed to ${url}`));
      });
      ws.addEventListener('close', () => {
        this.connections.delete(peerAddress);
      });
    });
  }

  private async request(
    ws: WebSocket,
    requestId: Uint8Array,
    frame: Uint8Array,
    signal?: AbortSignal,
  ): Promise<ArrayBuffer> {
    if (signal?.aborted) throw new TransferAbortError('aborted');
    return new Promise<ArrayBuffer>((resolve, reject) => {
      const onMsg = (ev: MessageEvent): void => {
        if (!(ev.data instanceof ArrayBuffer)) return;
        const buf = ev.data;
        if (buf.byteLength < 17) return;
        const idView = new Uint8Array(buf, 1, 16);
        if (!byteEqual(idView, requestId)) return;
        cleanup();
        resolve(buf);
      };
      const onErr = (): void => {
        cleanup();
        reject(new TransferTransportError('WS errored mid-request'));
      };
      const onClose = (): void => {
        cleanup();
        reject(new TransferTransportError('WS closed mid-request'));
      };
      const onAbort = (): void => {
        cleanup();
        reject(new TransferAbortError('aborted'));
      };
      function cleanup() {
        ws.removeEventListener('message', onMsg);
        ws.removeEventListener('error', onErr);
        ws.removeEventListener('close', onClose);
        signal?.removeEventListener('abort', onAbort);
      }
      ws.addEventListener('message', onMsg);
      ws.addEventListener('error', onErr);
      ws.addEventListener('close', onClose);
      signal?.addEventListener('abort', onAbort, { once: true });
      ws.send(frame as unknown as ArrayBuffer);
    });
  }
}

/**
 * Tries the primary transport first, falls back to secondary on
 * `TransferTransportError`. Does NOT retry per-chunk between transports —
 * the engine's retry layer handles that — but switches sticky after the
 * first failure so subsequent calls go to the fallback.
 */
export class FallbackTransferTransport implements ITransferTransport {
  private failed = false;

  constructor(
    private readonly primary: ITransferTransport,
    private readonly fallback: ITransferTransport,
  ) {}

  private active(): ITransferTransport {
    return this.failed ? this.fallback : this.primary;
  }

  async probe(peerAddress: string): Promise<void> {
    if (this.failed) return this.fallback.probe(peerAddress);
    try {
      await this.primary.probe(peerAddress);
    } catch (err) {
      if (err instanceof TransferTransportError) {
        this.failed = true;
        return this.fallback.probe(peerAddress);
      }
      throw err;
    }
  }

  async sendChunk(
    peerAddress: string,
    streamId: string,
    laneId: number,
    seq: number | bigint,
    bytes: Uint8Array,
    options?: ChunkSendOptions,
  ): Promise<ChunkAck> {
    if (this.failed) return this.fallback.sendChunk(peerAddress, streamId, laneId, seq, bytes, options);
    try {
      return await this.primary.sendChunk(peerAddress, streamId, laneId, seq, bytes, options);
    } catch (err) {
      if (err instanceof TransferTransportError) {
        this.failed = true;
        return this.fallback.sendChunk(peerAddress, streamId, laneId, seq, bytes, options);
      }
      throw err;
    }
  }

  async fetchResumeState(
    peerAddress: string,
    streamId: string,
  ): Promise<TransferResumeState | null> {
    return this.active().fetchResumeState(peerAddress, streamId);
  }

  /** True if the primary failed and we've fallen back. */
  get fellBack(): boolean {
    return this.failed;
  }
}

// ─── Helpers ─────────────────────────────────────────────────

function randomBytes(n: number): Uint8Array {
  const b = new Uint8Array(n);
  globalThis.crypto.getRandomValues(b);
  return b;
}

async function sha256(data: Uint8Array): Promise<Uint8Array> {
  const buf = await globalThis.crypto.subtle.digest('SHA-256', data as unknown as ArrayBuffer);
  return new Uint8Array(buf);
}

function byteEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) return false;
  }
  return true;
}

function base64UrlToBytes(s: string): Uint8Array {
  const padded = s.replace(/-/g, '+').replace(/_/g, '/');
  const pad = padded.length % 4 === 0 ? '' : '='.repeat(4 - (padded.length % 4));
  const bin = atob(padded + pad);
  const out = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
  return out;
}
