// WarzonePhone — WZP-WS-FEC client (Variant 5).
// WebSocket transport, WZP wire protocol, WASM RaptorQ FEC.
// Application-layer redundancy even over TCP.
// Sends MediaPacket-formatted frames with FEC encoding.
// Ready for direct relay WS support (no bridge translation needed).

'use strict';

// WASM module path (served from /wasm/ by the wzp-web bridge).
const WZP_WS_FEC_WASM_PATH = '/wasm/wzp_wasm.js';

// 12-byte MediaHeader size (matches wzp-proto MediaHeader::WIRE_SIZE).
const WZP_WS_FEC_HEADER_SIZE = 12;

// FEC wire header: block_id(1) + symbol_idx(1) + is_repair(1) = 3 bytes.
const WZP_WS_FEC_FEC_HEADER_SIZE = 3;

// FEC parameters.
// A 960-sample Int16 PCM frame = 1920 bytes. We use symbol_size = 2048
// (1920 payload + 2-byte length prefix + 126 bytes padding).
const WZP_WS_FEC_BLOCK_SIZE = 5;
const WZP_WS_FEC_SYMBOL_SIZE = 2048;

// Length prefix size within each FEC symbol.
const WZP_WS_FEC_LENGTH_PREFIX = 2;

class WZPWsFecClient {
  /**
   * @param {Object} options
   * @param {string} options.wsUrl       WebSocket URL (ws://host/ws/room)
   * @param {string} options.room        Room name
   * @param {Function} options.onAudio   callback(Int16Array) for playback
   * @param {Function} options.onStatus  callback(string) for UI status
   * @param {Function} options.onStats   callback(Object) for UI stats
   */
  constructor(options) {
    this.wsUrl = options.wsUrl;
    this.room = options.room;
    this.authToken = options.authToken || null;
    this.onAudio = options.onAudio || null;
    this.onStatus = options.onStatus || null;
    this.onStats = options.onStats || null;

    this.ws = null;
    this.seq = 0;
    this.startTimestamp = 0;
    this.stats = { sent: 0, recv: 0, fecRecovered: 0 };
    this._startTime = 0;
    this._statsInterval = null;
    this._connected = false;
    this._authenticated = false;

    // WASM FEC instances (loaded in loadWasm() / connect()).
    this._wasmModule = null;
    this.fecEncoder = null;
    this.fecDecoder = null;
    this.wasmReady = false;

    // Current FEC block counter for outgoing packets.
    this._fecBlockId = 0;
  }

  /**
   * Load the WASM FEC module.
   * Called automatically by connect(), or can be called early.
   */
  async loadWasm() {
    if (this.wasmReady) return;
    try {
      this._wasmModule = await import(WZP_WS_FEC_WASM_PATH);
      await this._wasmModule.default();

      this.fecEncoder = new this._wasmModule.WzpFecEncoder(
        WZP_WS_FEC_BLOCK_SIZE,
        WZP_WS_FEC_SYMBOL_SIZE
      );
      this.fecDecoder = new this._wasmModule.WzpFecDecoder(
        WZP_WS_FEC_BLOCK_SIZE,
        WZP_WS_FEC_SYMBOL_SIZE
      );
      this.wasmReady = true;
      console.log('[wzp-ws-fec] WASM FEC module loaded successfully');
    } catch (e) {
      console.error('[wzp-ws-fec] WASM FEC module failed to load:', e);
      this.wasmReady = false;
      throw e;
    }
  }

  /**
   * Build a 12-byte WZP MediaHeader.
   *
   * @param {number} seq         Sequence number (u16)
   * @param {number} timestampMs Milliseconds since session start
   * @param {boolean} isRepair   True if this is a FEC repair symbol
   * @param {number} codecId     Codec ID (0=RawPcm16, 1=Opus16k, 2=Opus48k)
   * @param {number} fecBlock    FEC block ID (u8)
   * @param {number} fecSymbol   FEC symbol index (u8)
   * @param {number} fecRatio    FEC ratio (0.0 to ~2.0)
   * @param {boolean} hasQuality Whether a quality report is attached
   * @returns {Uint8Array} 12-byte header
   */
  _buildHeader(seq, timestampMs, isRepair = false, codecId = 0, fecBlock = 0, fecSymbol = 0, fecRatio = 0, hasQuality = false) {
    const buf = new ArrayBuffer(WZP_WS_FEC_HEADER_SIZE);
    const view = new DataView(buf);

    const fecRatioEncoded = Math.min(127, Math.round(fecRatio * 63.5));
    const byte0 = ((0 & 0x01) << 7)           // version=0
        | ((isRepair ? 1 : 0) << 6)            // T bit
        | ((codecId & 0x0F) << 2)              // CodecID
        | ((hasQuality ? 1 : 0) << 1)          // Q bit
        | ((fecRatioEncoded >> 6) & 0x01);     // FecRatioHi
    view.setUint8(0, byte0);

    const byte1 = (fecRatioEncoded & 0x3F) << 2;
    view.setUint8(1, byte1);

    view.setUint16(2, seq & 0xFFFF);           // big-endian (default for DataView)
    view.setUint32(4, timestampMs & 0xFFFFFFFF); // big-endian
    view.setUint8(8, fecBlock & 0xFF);
    view.setUint8(9, fecSymbol & 0xFF);
    view.setUint8(10, 0);  // reserved
    view.setUint8(11, 0);  // csrc_count
    return new Uint8Array(buf);
  }

  /**
   * Parse a 12-byte MediaHeader from received binary data.
   *
   * @param {Uint8Array} data  At least 12 bytes
   * @returns {Object|null} Parsed header fields, or null if too short
   */
  _parseHeader(data) {
    if (data.byteLength < WZP_WS_FEC_HEADER_SIZE) return null;
    const view = new DataView(data.buffer || data, data.byteOffset || 0, 12);
    const byte0 = view.getUint8(0);
    const byte1 = view.getUint8(1);
    const fecRatioEncoded = ((byte0 & 0x01) << 6) | ((byte1 >> 2) & 0x3F);
    return {
      version: (byte0 >> 7) & 1,
      isRepair: !!((byte0 >> 6) & 1),
      codecId: (byte0 >> 2) & 0x0F,
      hasQuality: !!((byte0 >> 1) & 1),
      fecRatio: fecRatioEncoded / 63.5,
      seq: view.getUint16(2),
      timestamp: view.getUint32(4),
      fecBlock: view.getUint8(8),
      fecSymbol: view.getUint8(9),
      reserved: view.getUint8(10),
      csrcCount: view.getUint8(11),
    };
  }

  /**
   * Pad a PCM frame into a FEC symbol with a 2-byte length prefix.
   * Symbol layout: [len_hi, len_lo, ...pcm_bytes..., ...zero_padding...]
   *
   * @param {Uint8Array} pcmBytes  Raw PCM bytes
   * @returns {Uint8Array} Padded symbol of WZP_WS_FEC_SYMBOL_SIZE bytes
   */
  _padToSymbol(pcmBytes) {
    const symbol = new Uint8Array(WZP_WS_FEC_SYMBOL_SIZE);
    const len = pcmBytes.length;
    symbol[0] = (len >> 8) & 0xFF;
    symbol[1] = len & 0xFF;
    symbol.set(pcmBytes, WZP_WS_FEC_LENGTH_PREFIX);
    return symbol;
  }

  /**
   * Extract the original PCM payload from a FEC symbol (strip prefix + padding).
   *
   * @param {Uint8Array} symbol  Symbol data (WZP_WS_FEC_SYMBOL_SIZE bytes)
   * @returns {Uint8Array} Original PCM bytes
   */
  _unpadSymbol(symbol) {
    const len = (symbol[0] << 8) | symbol[1];
    if (len > WZP_WS_FEC_SYMBOL_SIZE - WZP_WS_FEC_LENGTH_PREFIX) {
      // Sanity check: if length is bogus, return empty.
      return new Uint8Array(0);
    }
    return symbol.slice(WZP_WS_FEC_LENGTH_PREFIX, WZP_WS_FEC_LENGTH_PREFIX + len);
  }

  /**
   * Open WebSocket connection and load the WASM FEC module.
   * @returns {Promise<void>} resolves when connected
   */
  async connect() {
    if (this._connected) return;

    // Load WASM module in parallel with WebSocket connect.
    const wasmPromise = this.loadWasm();

    const wsPromise = new Promise((resolve, reject) => {
      this._status('Connecting (WZP-WS-FEC) to room: ' + this.room + '...');

      this.ws = new WebSocket(this.wsUrl);
      this.ws.binaryType = 'arraybuffer';

      this.ws.onopen = () => {
        // Send auth if token provided.
        if (this.authToken) {
          this.ws.send(JSON.stringify({ type: 'auth', token: this.authToken }));
        }

        this._connected = true;
        this._authenticated = !this.authToken;
        this.seq = 0;
        this.startTimestamp = Date.now();
        this.stats = { sent: 0, recv: 0, fecRecovered: 0 };
        this._startTime = Date.now();
        this._fecBlockId = 0;
        this._startStatsTimer();
        resolve();
      };

      this.ws.onmessage = (event) => {
        // Handle text messages (auth responses).
        if (typeof event.data === 'string') {
          try {
            const msg = JSON.parse(event.data);
            if (msg.type === 'auth_ok') {
              this._authenticated = true;
              this._status('Authenticated (WZP-WS-FEC) to room: ' + this.room);
            }
            if (msg.type === 'auth_error') {
              this._status('Auth failed: ' + (msg.reason || 'unknown'));
              this.disconnect();
            }
          } catch(e) { /* ignore non-JSON text */ }
          return;
        }
        this._handleMessage(event);
      };

      this.ws.onclose = () => {
        const was = this._connected;
        this._cleanup();
        if (was) this._status('Disconnected');
      };

      this.ws.onerror = () => {
        if (!this._connected) {
          this._cleanup();
          reject(new Error('WebSocket connection failed'));
        } else {
          this._status('Connection error');
        }
      };
    });

    await Promise.all([wasmPromise, wsPromise]);

    const fecStatus = this.wasmReady ? 'FEC ready' : 'FEC unavailable';
    this._status('Connected (WZP-WS-FEC) to room: ' + this.room + ' (' + fecStatus + ')');
  }

  /**
   * Close WebSocket and clean up.
   */
  disconnect() {
    this._connected = false;
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this._stopStatsTimer();
    // Keep WASM module loaded (reusable), but reset encoder/decoder.
    if (this.fecEncoder) {
      try { this.fecEncoder.free(); } catch (_) { /* ignore */ }
      this.fecEncoder = null;
    }
    if (this.fecDecoder) {
      try { this.fecDecoder.free(); } catch (_) { /* ignore */ }
      this.fecDecoder = null;
    }
  }

  /**
   * Send a PCM audio frame with FEC encoding over the WebSocket.
   *
   * Each PCM frame is padded to a FEC symbol (2048 bytes with length prefix)
   * and fed to the FEC encoder. When a block of 5 symbols completes, the
   * encoder outputs source + repair symbols. Each is sent as an individual
   * WZP MediaPacket with the appropriate fecBlock, fecSymbol, and isRepair
   * fields in the 12-byte header.
   *
   * @param {ArrayBuffer} pcmBuffer  960-sample Int16 PCM (1920 bytes)
   */
  async sendAudio(pcmBuffer) {
    if (!this._connected || !this.ws || this.ws.readyState !== WebSocket.OPEN) return;
    if (!this.wasmReady || !this.fecEncoder) return;

    const pcmBytes = new Uint8Array(pcmBuffer);

    // Pad PCM frame to FEC symbol size with length prefix.
    const symbol = this._padToSymbol(pcmBytes);

    // Feed to FEC encoder. Returns wire data when block completes.
    const fecOutput = this.fecEncoder.add_symbol(symbol);

    if (fecOutput) {
      // Block completed — send all packets (source + repair).
      const packetSize = WZP_WS_FEC_FEC_HEADER_SIZE + WZP_WS_FEC_SYMBOL_SIZE;
      const timestampMs = Date.now() - this.startTimestamp;

      for (let offset = 0; offset + packetSize <= fecOutput.length; offset += packetSize) {
        const blockId = fecOutput[offset];
        const symbolIdx = fecOutput[offset + 1];
        const isRepair = fecOutput[offset + 2] !== 0;
        const symbolData = fecOutput.slice(
          offset + WZP_WS_FEC_FEC_HEADER_SIZE,
          offset + packetSize
        );

        // Build WZP MediaHeader for this FEC symbol.
        // fecRatio ~0.5 for 50% repair overhead: encoded = round(0.5 * 63.5) = 32
        const header = this._buildHeader(
          this.seq,
          timestampMs,
          isRepair,
          0,          // codecId = RawPcm16
          blockId,
          symbolIdx,
          0.5,        // fecRatio
          false       // hasQuality
        );

        // Wire frame: header(12) + symbol_data(2048)
        const packet = new Uint8Array(WZP_WS_FEC_HEADER_SIZE + symbolData.length);
        packet.set(header, 0);
        packet.set(symbolData, WZP_WS_FEC_HEADER_SIZE);

        this.ws.send(packet.buffer);
        this.seq = (this.seq + 1) & 0xFFFF;
        this.stats.sent++;
      }

      this._fecBlockId++;
    }
    // If block not yet complete, accumulate (no packets sent yet).
  }

  /**
   * Test FEC encode -> simulate loss -> decode in the browser.
   * Demonstrates that the WASM RaptorQ module works correctly
   * with the WZP wire protocol symbol format.
   *
   * @param {Object} [opts]
   * @param {number} [opts.blockSize=5]     Source symbols per block
   * @param {number} [opts.symbolSize=2048] Padded symbol size
   * @param {number} [opts.frameSize=1920]  PCM frame size in bytes
   * @param {number} [opts.dropCount=2]     Number of packets to drop (simulated 30%+ loss)
   * @returns {Object} Test results
   */
  testFec(opts) {
    if (!this.wasmReady || !this._wasmModule) {
      return { success: false, error: 'WASM FEC module not loaded' };
    }

    const blockSize = (opts && opts.blockSize) || 5;
    const symbolSize = (opts && opts.symbolSize) || WZP_WS_FEC_SYMBOL_SIZE;
    const frameSize = (opts && opts.frameSize) || 1920;
    const dropCount = (opts && opts.dropCount) || 2;

    const FEC_HDR = 3; // block_id + symbol_idx + is_repair
    const packetSize = FEC_HDR + symbolSize;

    const t0 = performance.now();

    // Create fresh encoder/decoder for the test.
    const encoder = new this._wasmModule.WzpFecEncoder(blockSize, symbolSize);
    const decoder = new this._wasmModule.WzpFecDecoder(blockSize, symbolSize);

    // Generate test frames with known data, padded to symbol size with length prefix.
    const originalFrames = [];
    const paddedSymbols = [];
    for (let i = 0; i < blockSize; i++) {
      const frame = new Uint8Array(frameSize);
      for (let j = 0; j < frameSize; j++) {
        frame[j] = ((i * 37 + 7) + j) & 0xFF;
      }
      originalFrames.push(frame);

      // Pad with length prefix (same as _padToSymbol).
      const sym = new Uint8Array(symbolSize);
      sym[0] = (frameSize >> 8) & 0xFF;
      sym[1] = frameSize & 0xFF;
      sym.set(frame, 2);
      paddedSymbols.push(sym);
    }

    // Encode: feed padded symbols to encoder.
    let wireData = null;
    for (const sym of paddedSymbols) {
      const result = encoder.add_symbol(sym);
      if (result) wireData = result;
    }

    if (!wireData) {
      wireData = encoder.flush();
    }

    // Parse wire packets.
    const packets = [];
    if (wireData) {
      for (let offset = 0; offset + packetSize <= wireData.length; offset += packetSize) {
        packets.push({
          blockId: wireData[offset],
          symbolIdx: wireData[offset + 1],
          isRepair: wireData[offset + 2] !== 0,
          data: wireData.slice(offset + FEC_HDR, offset + packetSize),
        });
      }
    }

    const sourcePackets = packets.filter(p => !p.isRepair).length;
    const repairPackets = packets.filter(p => p.isRepair).length;

    // Simulate packet loss: drop `dropCount` source packets from the front.
    const dropped = [];
    const surviving = [];
    for (let i = 0; i < packets.length; i++) {
      if (i < dropCount) {
        dropped.push(i);
      } else {
        surviving.push(packets[i]);
      }
    }

    // Decode from surviving packets.
    let decoded = null;
    for (const pkt of surviving) {
      const result = decoder.add_symbol(pkt.blockId, pkt.symbolIdx, pkt.isRepair, pkt.data);
      if (result) {
        decoded = result;
        break;
      }
    }

    // Verify decoded data: extract original frames from decoded symbols.
    let success = false;
    if (decoded) {
      // decoded is the concatenated padded symbols. Extract original frames.
      const recoveredFrames = [];
      for (let i = 0; i < blockSize; i++) {
        const symOffset = i * symbolSize;
        if (symOffset + symbolSize <= decoded.length) {
          const sym = decoded.slice(symOffset, symOffset + symbolSize);
          const len = (sym[0] << 8) | sym[1];
          recoveredFrames.push(sym.slice(2, 2 + len));
        }
      }

      success = recoveredFrames.length === blockSize;
      if (success) {
        for (let i = 0; i < blockSize && success; i++) {
          if (recoveredFrames[i].length !== originalFrames[i].length) {
            success = false;
            break;
          }
          for (let j = 0; j < originalFrames[i].length; j++) {
            if (recoveredFrames[i][j] !== originalFrames[i][j]) {
              success = false;
              break;
            }
          }
        }
      }
    }

    // Free WASM objects.
    encoder.free();
    decoder.free();

    const elapsed = performance.now() - t0;

    return {
      success,
      sourcePackets,
      repairPackets,
      totalPackets: packets.length,
      dropped: dropCount,
      recovered: !!decoded,
      symbolSize: symbolSize,
      frameSize: frameSize,
      elapsed: elapsed.toFixed(2) + 'ms',
    };
  }

  // -----------------------------------------------------------------------
  // Internal
  // -----------------------------------------------------------------------

  _handleMessage(event) {
    if (!(event.data instanceof ArrayBuffer)) return;
    const data = new Uint8Array(event.data);
    if (data.length < WZP_WS_FEC_HEADER_SIZE) return;

    const header = this._parseHeader(data);
    if (!header) return;

    this.stats.recv++;

    if (!this.wasmReady || !this.fecDecoder) {
      // No FEC decoder — cannot process FEC-encoded data.
      return;
    }

    // Extract symbol data (everything after 12-byte MediaHeader).
    const symbolData = data.slice(WZP_WS_FEC_HEADER_SIZE);

    // Feed symbol to FEC decoder using header fields.
    const decoded = this.fecDecoder.add_symbol(
      header.fecBlock,
      header.fecSymbol,
      header.isRepair,
      symbolData
    );

    if (decoded) {
      this.stats.fecRecovered++;

      // decoded is concatenated padded symbols.
      // Each symbol is WZP_WS_FEC_SYMBOL_SIZE bytes with a 2-byte length prefix.
      for (let off = 0; off + WZP_WS_FEC_SYMBOL_SIZE <= decoded.length; off += WZP_WS_FEC_SYMBOL_SIZE) {
        const symbol = decoded.slice(off, off + WZP_WS_FEC_SYMBOL_SIZE);
        const pcmBytes = this._unpadSymbol(symbol);

        if (pcmBytes.length > 0 && pcmBytes.length % 2 === 0) {
          const pcm = new Int16Array(
            pcmBytes.buffer,
            pcmBytes.byteOffset,
            pcmBytes.byteLength / 2
          );
          if (this.onAudio) this.onAudio(pcm);
        }
      }
    }
  }

  _startStatsTimer() {
    this._stopStatsTimer();
    this._statsInterval = setInterval(() => {
      if (!this._connected) {
        this._stopStatsTimer();
        return;
      }
      const elapsed = (Date.now() - this._startTime) / 1000;
      const loss = this.stats.sent > 0
        ? Math.max(0, 1 - this.stats.recv / this.stats.sent)
        : 0;
      if (this.onStats) {
        this.onStats({
          sent: this.stats.sent,
          recv: this.stats.recv,
          loss: loss,
          elapsed: elapsed,
          fecRecovered: this.stats.fecRecovered,
          fecReady: this.wasmReady,
        });
      }
    }, 1000);
  }

  _stopStatsTimer() {
    if (this._statsInterval) {
      clearInterval(this._statsInterval);
      this._statsInterval = null;
    }
  }

  _status(msg) {
    if (this.onStatus) this.onStatus(msg);
  }

  _cleanup() {
    this._connected = false;
    this._stopStatsTimer();
    if (this.ws) {
      try { this.ws.close(); } catch (_) { /* ignore */ }
      this.ws = null;
    }
  }
}

// ---------------------------------------------------------------------------
// Export
// ---------------------------------------------------------------------------

window.WZPWsFecClient = WZPWsFecClient;